Moved most main.rs functionality into lib.rs
This commit is contained in:
73
src/lib.rs
73
src/lib.rs
@@ -1,5 +1,12 @@
|
||||
// Robotnik libraries
|
||||
|
||||
use std::{os::unix::fs, sync::Arc};
|
||||
|
||||
use color_eyre::{Result, eyre::WrapErr};
|
||||
use human_panic::setup_panic;
|
||||
use tracing::{Level, info};
|
||||
use tracing_subscriber::FmtSubscriber;
|
||||
|
||||
pub mod chat;
|
||||
pub mod event;
|
||||
pub mod event_manager;
|
||||
@@ -10,3 +17,69 @@ 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);
|
||||
ev_manager_clone
|
||||
.broadcast(&Event::new("Starting..."))
|
||||
.await?;
|
||||
|
||||
let mut c = chat::new(&config, &handle).await?;
|
||||
|
||||
tokio::select! {
|
||||
_ = ev_manager_clone.start_listening("/tmp/robo.sock") => {
|
||||
// Event listener ended
|
||||
}
|
||||
result = c.run() => {
|
||||
result.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
77
src/main.rs
77
src/main.rs
@@ -1,79 +1,6 @@
|
||||
use color_eyre::{Result, eyre::WrapErr};
|
||||
use human_panic::setup_panic;
|
||||
use std::{os::unix::fs, sync::Arc};
|
||||
use tracing::{Level, info};
|
||||
use tracing_subscriber::FmtSubscriber;
|
||||
|
||||
use robotnik::{chat, Event, EventManager, qna, setup};
|
||||
|
||||
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.";
|
||||
use color_eyre::Result;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
// Some error sprucing.
|
||||
better_panic::install();
|
||||
setup_panic!();
|
||||
|
||||
let subscriber = FmtSubscriber::builder()
|
||||
.with_max_level(Level::TRACE)
|
||||
.finish();
|
||||
|
||||
tracing::subscriber::set_global_default(subscriber)
|
||||
.wrap_err("Failed to setup trace logging.")?;
|
||||
|
||||
info!("Starting");
|
||||
|
||||
let settings = setup::init().await.wrap_err("Failed to initialize.")?;
|
||||
let config = settings.config;
|
||||
|
||||
// chroot if applicable.
|
||||
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.")?;
|
||||
}
|
||||
|
||||
// Setup root path for commands.
|
||||
// let cmd_root = if let Ok(command_path) = config.get_string("command-path") {
|
||||
// Some(commands::Root::new(command_path))
|
||||
// } else {
|
||||
// None
|
||||
// };
|
||||
|
||||
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);
|
||||
ev_manager_clone
|
||||
.broadcast(&Event::new("Starting..."))
|
||||
.await?;
|
||||
|
||||
let mut c = chat::new(&config, &handle).await?;
|
||||
|
||||
tokio::select! {
|
||||
_ = ev_manager_clone.start_listening("/tmp/robo.sock") => {
|
||||
// Event listener ended
|
||||
}
|
||||
result = c.run() => {
|
||||
result.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
robotnik::run().await
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user