doubleo7/src/main.rs

use std::sync::LazyLock;
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]
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.
    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)
}