rssto/crates/llm/src/main.rs

mod argument;
mod config;

use anyhow::Result;
use mysql::Database;

#[tokio::main]
async fn main() -> Result<()> {
    use chrono::Local;
    use clap::Parser;
    use lancor::{ChatCompletionRequest, LlamaCppClient, Message};
    use log::{debug, info};

    use std::env::var;

    if var("RUST_LOG").is_ok() {
        use tracing_subscriber::{EnvFilter, fmt::*};
        struct T;
        impl time::FormatTime for T {
            fn format_time(&self, w: &mut format::Writer<'_>) -> std::fmt::Result {
                write!(w, "{}", Local::now())
            }
        }
        fmt()
            .with_timer(T)
            .with_env_filter(EnvFilter::from_default_env())
            .init()
    }

    let argument = argument::Argument::parse();
    let config: config::Config = toml::from_str(&std::fs::read_to_string(argument.config)?)?;

    let llm = LlamaCppClient::new(format!(
        "{}://{}:{}",
        config.llm.scheme, config.llm.host, config.llm.port
    ))?;
    let db = Database::pool(
        &config.mysql.host.to_string(),
        config.mysql.port,
        &config.mysql.username,
        &config.mysql.password,
        &config.mysql.database,
    )?;

    let provider_id = {
        let mut conn = db.connection()?;
        match conn.provider_id_by_name(&config.llm.model)? {
            Some(provider_id) => {
                debug!(
                    "Use existing DB provider #{} matches model name `{}`",
                    provider_id, &config.llm.model
                );
                provider_id
            }
            None => {
                let provider_id = conn.insert_provider(&config.llm.model)?;
                info!(
                    "Provider `{}` not found in database, created new one with ID `{provider_id}`",
                    &config.llm.model
                );
                provider_id
            }
        }
    };

    info!("Daemon started");
    loop {
        debug!("New queue begin...");
        let mut tx = db.transaction()?;
        for source in tx.contents_queue_for_provider_id(provider_id)? {
            debug!(
                "Begin generating `content_id` #{} using `provider_id` #{provider_id}.",
                source.content_id
            );

            let title = llm
                .chat_completion(ChatCompletionRequest::new(&config.llm.model).message(
                    Message::user(format!("{}\n{}", config.llm.message, source.title)),
                ))
                .await?;

            let description = llm
                .chat_completion(ChatCompletionRequest::new(&config.llm.model).message(
                    Message::user(format!("{}\n{}", config.llm.message, source.description)),
                ))
                .await?;

            let content_id = tx.insert_content(
                source.channel_item_id,
                Some(provider_id),
                &title.choices[0].message.content,
                &description.choices[0].message.content,
            )?;

            debug!(
                "Created `content_id` #{content_id} using `content_id` #{} source by `provider_id` #{provider_id}.",
                source.content_id
            )
        }
        tx.commit()?;
        debug!("Queue completed");
        if let Some(update) = config.update {
            debug!("Wait {update} seconds to continue...");
            std::thread::sleep(std::time::Duration::from_secs(update))
        } else {
            return Ok(());
        }
    }
}