doubleo7/swear_cleanup/src/starter.rs

use anyhow;
use rig::client::{CompletionClient, Nothing};
use rig::providers::llamafile::LlamafileExt;
use rig::providers::openai::GenericCompletionModel;
use rig::providers::{llamafile, ollama};
use crate::{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)
}