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<()> {
    // 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?;

    println!("Final text (score={:.6}):\n{}", result.score, 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)
}