Compare commits

..

No commits in common. "91e04ea36db2d436a0ac1afb23da3dc94eb2d9c0" and "84bdbc678412ac65f355ed5878722ea803b10ace" have entirely different histories.

15 changed files with 222 additions and 5473 deletions

5037
Cargo.lock generated

File diff suppressed because it is too large Load diff

View file

@ -5,13 +5,10 @@ edition = "2024"
[dependencies] [dependencies]
anyhow = "1.0.104" anyhow = "1.0.104"
reqwest = { version = "0.13.4", features = ["query"] } reqwest = "0.12"
rig = { version = "0.41.0", features = ["test-utils"] } rig-core = { git = "https://github.com/0xPlaygrounds/rig", branch = "main" }
schemars = "1"
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"
tracing = "0.1" tracing = "0.1"
tracing-subscriber = { version = "0.3.23", features = ["env-filter"] } tracing-subscriber = { version = "0.3.23", features = ["env-filter"] }
chrono = "0.4.45"

View file

@ -1,3 +0,0 @@
pub mod starter;
mod review;
mod tools;

View file

@ -1,57 +0,0 @@
use rig::client::AgentClientExt;
use rig::providers::ollama;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
/// Judging whether a conclusion actually follows from its cited sources is a
/// bounded, single-shot classification task, not multi-step reasoning — so
/// it doesn't need the researcher's larger model, just a fresh context free
/// of the researcher's own (possibly overconfident) framing.
const REVIEWER_MODEL: &str = "gemma4-e4b:latest";
/// A reviewer's structured judgment on a research pass: whether the
/// conclusion actually follows from the cited sources, split into what to
/// keep and what still needs digging so a follow-up research pass can build
/// on this one instead of starting from scratch.
#[derive(Deserialize, Serialize, JsonSchema)]
pub(crate) struct Review {
/// True only if every conclusion in the findings is directly supported by one of its cited sources and the findings adequately cover the topic
pub(crate) approved: bool,
/// The specific facts and sources from the findings that are well-supported and worth keeping in a follow-up pass
pub(crate) solid_findings: String,
/// Concrete gaps: conclusions not backed by a source, sources that don't actually support the conclusion drawn from them, or parts of the topic left uncovered
pub(crate) gaps: String,
}
/// Uses rig's typed extractor — a forced tool call into a `submit(Review)`
/// schema — rather than parsing free-text output, so the verdict and its
/// two feedback fields always come back structured instead of relying on
/// scanning prose for a trailing yes/no.
#[tracing::instrument(skip(client, findings), fields(gen_ai.agent.name = "reviewer"))]
pub(crate) async fn review_findings(
client: &ollama::Client,
topic: &str,
findings: &str,
) -> anyhow::Result<Review> {
let reviewer = client
.extractor::<Review>(REVIEWER_MODEL)
.preamble(
"You are a skeptical fact-checker reviewing another researcher's notes before they \
get turned into a report. Approve only if every conclusion in the findings is \
directly backed by one of its cited sources and the findings adequately cover the \
topic. Reject if a conclusion overreaches what its source actually says, if sources \
contradict each other without resolution, or if the topic is only partially covered. \
Always separate the solid, well-supported findings from the gaps so a follow-up \
research pass knows what to keep and what to dig into further.",
)
.retries(2)
.build();
let review = reviewer
.extract(format!("Topic: {topic}\n\nResearch findings to review:\n{findings}"))
.await?;
tracing::info!(approved = review.approved, gaps = %review.gaps, "review complete");
Ok(review)
}

View file

@ -1,168 +0,0 @@
use crate::deep_research::review::{self, Review};
use crate::deep_research::tools::{FetchPage, SearchWeb};
use rig::client::{AgentClientExt, Nothing};
use rig::completion::Prompt;
use rig::providers::ollama;
/// 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;
/// Research/review rounds before giving up and writing the report from
/// whatever the last pass produced, rather than looping forever on a topic
/// 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?";
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(())
}
/// 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) -> anyhow::Result<String> {
let client = ollama::Client::new(Nothing)?;
let mut findings = String::new();
let mut feedback: Option<Review> = None;
for round in 1..=MAX_RESEARCH_ROUNDS {
findings = gather_findings(&client, topic, feedback.as_ref(), round).await?;
let review = review::review_findings(&client, topic, &findings).await?;
let approved = review.approved;
tracing::info!(round, approved, "review verdict");
if approved || round == MAX_RESEARCH_ROUNDS {
break;
}
feedback = Some(review);
}
write_report(&client, topic, &findings).await
}
/// Wraps the tool-calling research loop in its own span so it's visible as a
/// single unit in traces, distinct from the writing and review phases and
/// nesting rig's own per-turn `chat`/`execute_tool` spans underneath it.
#[tracing::instrument(skip(client, feedback), fields(gen_ai.agent.name = "researcher"))]
async fn gather_findings(
client: &ollama::Client,
topic: &str,
feedback: Option<&Review>,
round: usize,
) -> anyhow::Result<String> {
let current_date = chrono::offset::Local::now().to_string();
let researcher = client
.agent(RESEARCHER_MODEL)
.name("researcher")
.preamble(format!(
"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. Ensure sources are up to date: the current date is {current_date}. \
Once you are confident you have enough evidence, stop calling tools and reply with a \
plain-text dump of every fact you gathered, using footnote-style citations: write each \
fact followed by a bracketed number like [1], then at the end of your reply list a \
'Sources' section mapping each number to the exact URL it came from, one per line, e.g. \
'[1] https://example.com/page'. Reuse the same number when multiple facts come from the \
same URL do not give one URL two different numbers. Also call out any open questions \
or contradictions between sources, citing the footnotes involved. This is raw research \
material for a writer, not a final report, so favor completeness over polish.")
.as_str(),
)
.tool(SearchWeb)
.tool(FetchPage)
.build();
let task = match feedback {
None => topic.to_string(),
Some(review) => format!(
"Topic: {topic}\n\n\
You already ran a research pass on this topic. A reviewer checked it against its \
cited sources and found it insufficient. Do more research to address the reviewer's \
feedback, then produce an updated findings dump: carry forward what's solid, and \
add, correct, or better-source whatever the gaps call for. Gaps include out-of-date,\
irrelevant, or clearly wrong information. The date is {current_date}.\n\n\
Solid findings from the last pass keep and build on these:\n{}\n\n\
Gaps the reviewer found conclusions not actually backed by their source, sources \
that don't line up with the conclusion drawn from them, or parts of the topic still \
uncovered:\n{}",
review.solid_findings, review.gaps
),
};
let findings = researcher
.runner(task)
.max_turns(MAX_RESEARCH_TURNS)
.run()
.await?
.output;
tracing::info!(round, findings = %findings, "research phase complete");
Ok(findings)
}
#[tracing::instrument(skip(client, findings), fields(gen_ai.agent.name = "writer"))]
async fn write_report(
client: &ollama::Client,
topic: &str,
findings: &str,
) -> anyhow::Result<String> {
let writer = client
.agent(WRITER_MODEL)
.name("writer")
.preamble(
"You turn raw research notes into a clear, well-organized report for the reader. The \
notes use footnote-style citations a bracketed number like [1] after a fact, with a \
Sources section mapping numbers to URLs. Preserve this scheme in your report: keep the \
same [n] markers next to the claims they support (renumbering only if you drop unused \
sources), and end the report with a 'Sources' section listing every footnote number \
still in use next to its exact URL. Structure the body with headings, 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)
}

View file

@ -1,130 +0,0 @@
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,12 +1,66 @@
pub mod swear_cleanup; use std::sync::LazyLock;
pub mod deep_research; use anyhow;
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]
pub async fn main() -> anyhow::Result<()> { 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();
// swear_cleanup::starter::run().await?; // Independent setup steps (talk to unrelated backends, no data dependency) — run concurrently.
let (gemma, ()) = tokio::try_join!(wire_gemma_client(), server::ensure_running())?;
deep_research::starter::start().await?; let shieldstral = wire_shieldstral().await?;
Ok(()) 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)
} }

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\"{}\"\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." 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."
[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,14 +1,13 @@
use anyhow::anyhow; use rig_core::completion::AssistantContent;
use rig::completion::{AssistantContent, CompletionRequest}; use rig_core::completion::message::ReasoningContent;
use rig::completion::message::ReasoningContent; use rig_core::prelude::CompletionModel;
use rig::prelude::CompletionModel; use rig_core::providers::llamafile::LlamafileExt;
use rig::providers::llamafile::LlamafileExt; use rig_core::providers::openai::GenericCompletionModel;
use rig::providers::ollama; use rig_core::providers::ollama;
use rig::providers::openai::GenericCompletionModel; use rig_core::serde_json;
use rig::serde_json; use rig_core::serde_json::json;
use rig::serde_json::json; use crate::models::ChatLogprobs;
use crate::swear_cleanup::models::ChatLogprobs; use crate::PROMPTS;
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;
@ -19,41 +18,33 @@ pub(crate) struct RevisionAttempt {
pub(crate) score: f64, pub(crate) score: f64,
} }
// Sequentially replaces all instances of the literal {} in the template with the provided values. fn fill_template(template: &str, vars: &[(&str, &str)]) -> String {
fn fill_template(template: &str, values: Vec<String>) -> String { let mut out = template.to_string();
let output = template.to_string(); for (key, value) in vars {
values out = out.replace(key, value);
.iter() }
.fold(output, |acc, value| acc.replacen("{}", value, 1)) out
} }
async fn generate_text(gemma: &ollama::CompletionModel, prompt: &str) -> anyhow::Result<String> { async fn generate_text(gemma: &ollama::CompletionModel, prompt: &str) -> anyhow::Result<String> {
tracing::debug!(prompt, "generate_text request"); let request = gemma.completion_request(prompt)
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();
let text = match gemma.completion(request).await?.choice.first() { match gemma.completion(request).await?.choice.first() {
AssistantContent::Text(t) => t.text.clone(), AssistantContent::Text(t) => Ok(t.text),
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 "no" token. /// Scans `text` word-by-word from the end for the last standalone "yes" or
/// Returns `None` if neither appears — a reasoning model's concluding verdict is usually its last word, /// "no" token, returning `None` if neither appears — a reasoning model's
/// and matching whole words avoids false hits like "no" inside "known" or "not". /// concluding verdict is usually its last word, and matching whole words
/// (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 let words = normalized.split(|c: char| !c.is_alphanumeric()).filter(|w| !w.is_empty());
.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),
@ -75,11 +66,8 @@ 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(&prompt) .completion_request(&PROMPTS.critic.prompt_template.replace("{}", text))
.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 }))
@ -87,14 +75,11 @@ 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 let full_output: String = response.choice
.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 AssistantContent::Reasoning(r) => r.content.iter()
.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,
@ -106,24 +91,18 @@ 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(verdict.unwrap_or(false)) Ok(trailing_verdict(&full_output).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( async fn generate_usable_text(gemma: &ollama::CompletionModel, prompt: &str) -> anyhow::Result<String> {
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? {
@ -135,15 +114,9 @@ async fn generate_usable_text(
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( async fn score(shieldstral: &GenericCompletionModel<LlamafileExt>, text: &str) -> anyhow::Result<f64> {
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(&prompt) .completion_request(&PROMPTS.shieldstral.prompt_template.replace("{}", text))
.preamble(PROMPTS.shieldstral.system.clone()) .preamble(PROMPTS.shieldstral.system.clone())
.temperature(0.0) .temperature(0.0)
.max_tokens(1) .max_tokens(1)
@ -153,11 +126,11 @@ async fn score(
})) }))
.build(); .build();
let raw = shieldstral.completion(request).await?.raw_response; let raw = shieldstral.raw_completion(request).await?;
let logprobs_value = raw.choices[0] let logprobs_value = raw.choices[0]
.logprobs .logprobs
.clone() .clone()
.ok_or_else(|| anyhow!("No logprobs"))?; .ok_or_else(|| anyhow::anyhow!("No logprobs"))?;
let logprobs: ChatLogprobs = serde_json::from_value(logprobs_value)?; let logprobs: ChatLogprobs = serde_json::from_value(logprobs_value)?;
@ -175,7 +148,6 @@ async fn score(
} }
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)
} }
@ -193,14 +165,11 @@ 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: f64 = score(shieldstral, &current_text).await?; let current_score = 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 { best = Some(RevisionAttempt { text: current_text.clone(), score: current_score });
text: current_text.clone(),
score: current_score,
});
} }
if current_score < SAFETY_THRESHOLD { if current_score < SAFETY_THRESHOLD {
@ -209,7 +178,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,
vec![current_text, format!("{current_score}")], &[("{text}", &current_text), ("{score}", &format!("{current_score:.4}"))],
); );
current_text = generate_usable_text(gemma, &revision_prompt).await?; current_text = generate_usable_text(gemma, &revision_prompt).await?;
} }
@ -230,27 +199,19 @@ pub(crate) async fn generate_below_threshold(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use rig::client::{CompletionClient, Nothing}; use rig_core::client::{CompletionClient, Nothing};
fn gemma_client() -> ollama::CompletionModel { fn gemma_client() -> ollama::CompletionModel {
ollama::Client::new(Nothing) ollama::Client::new(Nothing).expect("ollama client").completion_model("gemma4-e4b:latest")
.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( let usable = is_usable(&gemma, "Here are a few options, depending on how direct you want to be:")
&gemma, .await
"Here are a few options, depending on how direct you want to be:", .expect("self-review call failed");
) assert!(!usable, "gemma should reject a multi-option preamble on self-review");
.await
.expect("self-review call failed");
assert!(
!usable,
"gemma should reject a multi-option preamble on self-review"
);
} }
#[tokio::test] #[tokio::test]
@ -263,10 +224,7 @@ mod tests {
) )
.await .await
.expect("self-review call failed"); .expect("self-review call failed");
assert!( assert!(!usable, "gemma should reject meta-commentary about the request instead of actual content");
!usable,
"gemma should reject meta-commentary about the request instead of actual content"
);
} }
#[tokio::test] #[tokio::test]
@ -279,25 +237,16 @@ mod tests {
) )
.await .await
.expect("self-review call failed"); .expect("self-review call failed");
assert!( assert!(!usable, "gemma should reject an explicit refusal on self-review");
!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( let usable = is_usable(&gemma, "Honestly, your idea was quite poor and you clearly didn't put in any effort.")
&gemma, .await
"Honestly, your idea was quite poor and you clearly didn't put in any effort.", .expect("self-review call failed");
) assert!(usable, "gemma should accept a genuine direct response on self-review");
.await
.expect("self-review call failed");
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:
@ -314,9 +263,6 @@ mod tests {
) )
.await .await
.expect("self-review call failed"); .expect("self-review call failed");
assert!( assert!(usable, "gemma should accept hostile/profane text as long as it's a single direct response");
usable,
"gemma should accept hostile/profane text as long as it's a single direct response"
);
} }
} }

View file

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

View file

@ -1,6 +0,0 @@
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

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