robotnik/src/lib.rs

// Robotnik libraries

use std::{os::unix::fs, sync::Arc};

use color_eyre::{Result, eyre::WrapErr};
use human_panic::setup_panic;
use tokio::sync::mpsc;
use tracing::{Level, info};
use tracing_subscriber::FmtSubscriber;

pub mod chat;
pub mod command;
pub mod event;
pub mod event_manager;
pub mod plugin;
pub mod qna;
pub mod setup;

pub use event::Event;
pub use event_manager::EventManager;
pub use qna::LLMHandle;

const DEFAULT_INSTRUCT: &str =
    "You are a shady, yet helpful IRC bot.  You try to give responses that can
 be sent in a single IRC response according to the specification.  Keep answers to
 500 characters or less.";

// NB: Everything should fail if logging doesn't start properly.
async fn init_logging() {
    better_panic::install();
    setup_panic!();

    let subscriber = FmtSubscriber::builder()
        .with_max_level(Level::TRACE)
        .finish();

    tracing::subscriber::set_global_default(subscriber).unwrap();
}

pub async fn run() -> Result<()> {
    init_logging().await;
    info!("Starting up.");

    let settings = setup::init().await.wrap_err("Failed to initialize.")?;
    let config = settings.config;

    // NOTE: Doing chroot this way might be impractical.
    if let Ok(chroot_path) = config.get_string("chroot-dir") {
        info!("Attempting to chroot to {}", chroot_path);
        fs::chroot(&chroot_path)
            .wrap_err_with(|| format!("Failed setting chroot '{}'", chroot_path))?;
        std::env::set_current_dir("/").wrap_err("Couldn't change directory after chroot.")?;
    }

    let handle = qna::LLMHandle::new(
        config.get_string("api-key").wrap_err("API missing.")?,
        config
            .get_string("base-url")
            .wrap_err("base-url missing.")?,
        config
            .get_string("model")
            .wrap_err("model string missing.")?,
        config
            .get_string("instruct")
            .unwrap_or_else(|_| DEFAULT_INSTRUCT.to_string()),
    )
    .wrap_err("Couldn't initialize LLM handle.")?;

    let ev_manager = Arc::new(EventManager::new()?);
    let ev_manager_clone = Arc::clone(&ev_manager);

    let mut c = chat::new(&config, &handle, Arc::clone(&ev_manager)).await?;

    let (from_plugins, to_chat) = mpsc::channel(100);

    tokio::select! {
        _ = ev_manager_clone.start_listening("/tmp/robo.sock") => {
            // Event listener ended
        }
        result = c.run(to_chat) => {
            if let Err(e) = result {
                tracing::error!("Chat run error: {:?}", e);
                return Err(e);
            }
        }
        fifo = EventManager::start_fifo("/tmp/robo_in.sock", from_plugins) => {
            fifo.wrap_err("FIFO reader failed.")?;
        }
    }

    Ok(())
}