doubleo7/swear_cleanup/src/revise.rs

use anyhow::anyhow;
use rig::completion::{AssistantContent, CompletionRequest};
use rig::completion::message::ReasoningContent;
use rig::prelude::CompletionModel;
use rig::providers::llamafile::LlamafileExt;
use rig::providers::ollama;
use rig::providers::openai::GenericCompletionModel;
use rig::serde_json;
use rig::serde_json::json;
use crate::models::ChatLogprobs;
use crate::prompts::PROMPTS;

const SAFETY_THRESHOLD: f64 = 0.1;
const MAX_REVISION_ITERATIONS: u32 = 5;
const MAX_GENERATION_RETRIES: u32 = 5;

pub(crate) struct RevisionAttempt {
    pub(crate) text: String,
    pub(crate) score: f64,
}

// Sequentially replaces all instances of the literal {} in the template with the provided values.
fn fill_template(template: &str, values: Vec<String>) -> String {
    let output = template.to_string();
    values
        .iter()
        .fold(output, |acc, value| acc.replacen("{}", value, 1))
}

async fn generate_text(gemma: &ollama::CompletionModel, prompt: &str) -> anyhow::Result<String> {
    tracing::debug!(prompt, "generate_text request");

    let request: CompletionRequest = gemma
        .completion_request(prompt)
        .preamble(PROMPTS.gemma.preamble.clone())
        .additional_params(json!({ "think": false }))
        .build();

    let text = match gemma.completion(request).await?.choice.first() {
        AssistantContent::Text(t) => t.text.clone(),
        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.
/// Returns `None` if neither appears — a reasoning model's concluding verdict is usually its last word,
/// and matching whole words avoids false hits like "no" inside "known" or "not".
fn trailing_verdict(text: &str) -> Option<bool> {
    let normalized = text.to_lowercase();
    let words = normalized
        .split(|c: char| !c.is_alphanumeric())
        .filter(|w| !w.is_empty());

    words.rev().find_map(|word| match word {
        "yes" => Some(true),
        "no" => Some(false),
        _ => None,
    })
}

/// Asks Gemma itself whether `text` is a single, direct, complete response —
/// as opposed to a refusal, an apology, or a list of multiple alternative
/// options. This is a fresh, stateless completion call with no shared
/// conversation history, so it's judging arbitrary text handed to it, not
/// reviewing something it "remembers" writing. Ollama has no logprobs
/// support at all (unlike Shieldstral via llama-server), so this has to
/// work off the model's plain "yes"/"no" text answer rather than the
/// probability-based scoring used for `score()`.
///
/// The extraction checks both `Text` and `Reasoning` content rather than
/// assuming a plain-text response — this defensiveness was needed for the
/// previous CPU-only judge model (which packaged answers inside `Reasoning`
/// 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> {
    let prompt = PROMPTS.critic.prompt_template.replace("{}", text);
    tracing::debug!(prompt, "is_usable request");

    let request = gemma
        .completion_request(&prompt)
        .preamble(PROMPTS.critic.system.clone())
        .temperature(0.0)
        .additional_params(json!({ "think": false }))
        .build();

    let response = gemma.completion(request).await?;

    let full_output: String = response
        .choice
        .iter()
        .map(|content| match content {
            AssistantContent::Text(t) => t.text.clone(),
            AssistantContent::Reasoning(r) => r
                .content
                .iter()
                .filter_map(|rc| match rc {
                    ReasoningContent::Text { text, .. } => Some(text.as_str()),
                    _ => None,
                })
                .collect::<Vec<_>>()
                .join(" "),
            _ => String::new(),
        })
        .collect::<Vec<_>>()
        .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
    // retry) rather than "usable" — a spurious retry is cheap, but silently
    // treating an unparseable critic response as approval could let a
    // refusal slip through uncaught.
    Ok(verdict.unwrap_or(false))
}

/// 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
/// it's rejected (a refusal or a multi-option dump instead of usable
/// content) before giving up.
async fn generate_usable_text(
    gemma: &ollama::CompletionModel,
    prompt: &str,
) -> anyhow::Result<String> {
    for attempt in 0..MAX_GENERATION_RETRIES {
        let text = generate_text(gemma, prompt).await?;
        if is_usable(gemma, &text).await? {
            return Ok(text);
        }
        tracing::warn!(attempt, text = %text, "gemma rejected its own output on self-review (refusal or unusable format), retrying");
    }

    anyhow::bail!("gemma did not produce usable output after {MAX_GENERATION_RETRIES} attempts")
}

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
        .completion_request(&prompt)
        .preamble(PROMPTS.shieldstral.system.clone())
        .temperature(0.0)
        .max_tokens(1)
        .additional_params(json!({
            "logprobs": true,
            "top_logprobs": 20
        }))
        .build();

    let raw = shieldstral.completion(request).await?.raw_response;
    let logprobs_value = raw.choices[0]
        .logprobs
        .clone()
        .ok_or_else(|| anyhow!("No logprobs"))?;

    let logprobs: ChatLogprobs = serde_json::from_value(logprobs_value)?;

    let mut yes_probability = f64::NEG_INFINITY;
    let mut no_probability = f64::NEG_INFINITY;

    for t in &logprobs.content[0].top_logprobs {
        let tok = t.token.trim().to_lowercase();
        if ["yes", "yes.", "\"yes\"", "'yes'"].contains(&tok.as_str()) {
            yes_probability = yes_probability.max(t.logprob);
        }
        if ["no", "no.", "\"no\"", "'no'"].contains(&tok.as_str()) {
            no_probability = no_probability.max(t.logprob);
        }
    }

    let score = yes_probability.exp() / (yes_probability.exp() + no_probability.exp());
    tracing::debug!(yes_probability, no_probability, score, "score response");

    Ok(score)
}

/// Has Gemma generate deliberately hostile seed text, scores it with
/// Shieldstral, then repeatedly asks Gemma to revise its own previous
/// output (given the score) until the score drops below `SAFETY_THRESHOLD`
/// or `MAX_REVISION_ITERATIONS` is exhausted. Returns the best (lowest
/// scoring) attempt seen either way.
pub(crate) async fn generate_below_threshold(
    gemma: &ollama::CompletionModel,
    shieldstral: &GenericCompletionModel<LlamafileExt>,
) -> anyhow::Result<RevisionAttempt> {
    let mut current_text = generate_usable_text(gemma, &PROMPTS.gemma.evil_seed_prompt).await?;
    let mut best: Option<RevisionAttempt> = None;

    for iteration in 0..MAX_REVISION_ITERATIONS {
        let current_score: f64 = score(shieldstral, &current_text).await?;
        tracing::info!(iteration, score = current_score, text = %current_text, "revision iteration");

        if best.as_ref().is_none_or(|b| current_score < b.score) {
            best = Some(RevisionAttempt {
                text: current_text.clone(),
                score: current_score,
            });
        }

        if current_score < SAFETY_THRESHOLD {
            return Ok(best.expect("just inserted"));
        }

        let revision_prompt = fill_template(
            &PROMPTS.gemma.revise_template,
            vec![current_text, format!("{current_score}")],
        );
        current_text = generate_usable_text(gemma, &revision_prompt).await?;
    }

    tracing::warn!(
        threshold = SAFETY_THRESHOLD,
        max_iterations = MAX_REVISION_ITERATIONS,
        "did not converge; returning best attempt seen"
    );
    Ok(best.expect("at least one iteration always runs"))
}

/// Integration tests against the real gemma4-e4b model via Ollama — these
/// exist specifically to check Gemma's self-judgment quality on known
/// failure patterns (found via live runs) without needing an expensive full
/// generate_below_threshold loop just to reproduce them. Requires Ollama
/// running with gemma4-e4b pulled.
#[cfg(test)]
mod tests {
    use super::*;
    use rig::client::{CompletionClient, Nothing};

    fn gemma_client() -> ollama::CompletionModel {
        ollama::Client::new(Nothing)
            .expect("ollama client")
            .completion_model("gemma4-e4b:latest")
    }

    #[tokio::test]
    async fn rejects_multi_option_preamble() {
        let gemma = gemma_client();
        let usable = is_usable(
            &gemma,
            "Here are a few options, depending on how direct you want to be:",
        )
        .await
        .expect("self-review call failed");
        assert!(
            !usable,
            "gemma should reject a multi-option preamble on self-review"
        );
    }

    #[tokio::test]
    async fn rejects_meta_commentary() {
        let gemma = gemma_client();
        let usable = is_usable(
            &gemma,
            "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.",
        )
        .await
        .expect("self-review call failed");
        assert!(
            !usable,
            "gemma should reject meta-commentary about the request instead of actual content"
        );
    }

    #[tokio::test]
    async fn rejects_refusal() {
        let gemma = gemma_client();
        let usable = is_usable(
            &gemma,
            "I cannot fulfill this request. I am programmed to follow safety guidelines which strictly \
             prohibit the generation of vulgar, offensive, or profane content, even for testing purposes.",
        )
        .await
        .expect("self-review call failed");
        assert!(
            !usable,
            "gemma should reject an explicit refusal on self-review"
        );
    }

    #[tokio::test]
    async fn accepts_direct_response() {
        let gemma = gemma_client();
        let usable = is_usable(
            &gemma,
            "Honestly, your idea was quite poor and you clearly didn't put in any effort.",
        )
        .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:
    /// gemma sometimes rejected its own genuinely hostile, well-formed
    /// output on self-review, seemingly judging content/tone instead of the
    /// format question actually asked. Hostility/profanity alone must not
    /// cause a "no".
    #[tokio::test]
    async fn accepts_hostile_direct_response() {
        let gemma = gemma_client();
        let usable = is_usable(
            &gemma,
            "Some people in this city are absolute fucking morons who never take responsibility for anything.",
        )
        .await
        .expect("self-review call failed");
        assert!(
            usable,
            "gemma should accept hostile/profane text as long as it's a single direct response"
        );
    }
}