145 lines
3.9 KiB
Rust
145 lines
3.9 KiB
Rust
use std::env;
|
|
use std::collections::HashSet;
|
|
use std::sync::Arc;
|
|
use serenity::http::Http;
|
|
use serenity::prelude::*;
|
|
use songbird::{SerenityInit, Songbird};
|
|
use reqwest::Client as HttpClient;
|
|
use serenity::all::{Cache, ShardManager, UserId};
|
|
use crate::api::App;
|
|
use crate::bot::handler::BotHandler;
|
|
use crate::bot::oai::OAI;
|
|
|
|
mod api;
|
|
mod bot;
|
|
mod data;
|
|
mod error;
|
|
mod utils;
|
|
|
|
pub struct HttpKey;
|
|
|
|
impl TypeMapKey for HttpKey {
|
|
type Value = HttpClient;
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct AppState {
|
|
client: reqwest::Client,
|
|
client_id: String,
|
|
client_secret: String,
|
|
redirect_uri: String,
|
|
http: Arc<Http>,
|
|
cache: Arc<Cache>,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
dotenv::dotenv().ok();
|
|
env_logger::init_from_env(env_logger::Env::default().filter_or("RUST_LOG", "warn,siren=info"));
|
|
|
|
data::initialize().await?;
|
|
|
|
let token: String = env::var("DISCORD_TOKEN").expect("Expected a token in the environment");
|
|
|
|
// Set up handler with optional OpenAI integration
|
|
let handler = configure_handler();
|
|
|
|
// Set up Songbird for voice functionality
|
|
let songbird = Songbird::serenity();
|
|
|
|
let intents: GatewayIntents = GatewayIntents::all();
|
|
|
|
let mut client = Client::builder(token, intents)
|
|
.event_handler(handler)
|
|
// .framework(StandardFramework::new().configure(|c| c.owners(owners)))
|
|
.register_songbird_with(Arc::clone(&songbird))
|
|
.type_map_insert::<HttpKey>(HttpClient::new())
|
|
.await
|
|
.expect("Error creating client");
|
|
|
|
let (bot_owner, bot_id) = get_bot_info(&client.http).await;
|
|
|
|
let client_secret: String =
|
|
env::var("DISCORD_SECRET").expect("Expected a secret in the environment");
|
|
let redirect_uri: String =
|
|
env::var("API_CALLBACK_URI").expect("Expected a secret in the environment");
|
|
let app_state = AppState {
|
|
client: HttpClient::new(),
|
|
client_id: bot_id.to_string(),
|
|
client_secret,
|
|
redirect_uri,
|
|
http: Arc::clone(&client.http),
|
|
cache: Arc::clone(&client.cache),
|
|
};
|
|
|
|
log::debug!("Starting Siren with ID: {bot_id} (Contact: {:?})", bot_owner);
|
|
|
|
// Spawn shutdown signal handling
|
|
let shard_manager = Arc::clone(&client.shard_manager);
|
|
tokio::spawn(async move {
|
|
signal_shutdown(shard_manager).await;
|
|
});
|
|
|
|
// Start API server
|
|
tokio::spawn(App::new(app_state).serve());
|
|
|
|
// Start Discord bot
|
|
if let Err(why) = client.start_autosharded().await {
|
|
log::error!("Client error: {why:?}");
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn get_bot_info(http: &Http) -> (Option<UserId>, UserId) {
|
|
match http.get_current_application_info().await {
|
|
Ok(info) => {
|
|
let bot_owner;
|
|
if let Some(team) = info.team {
|
|
bot_owner = Some(team.owner_user_id);
|
|
} else if let Some(owner) = info.owner {
|
|
bot_owner = Some(owner.id);
|
|
} else {
|
|
bot_owner = None;
|
|
}
|
|
match http.get_current_user().await {
|
|
Ok(bot) => (bot_owner, bot.id),
|
|
Err(why) => panic!("Could not access the bot id: {why:?}"),
|
|
}
|
|
}
|
|
Err(why) => panic!("Could not access application info: {why:?}"),
|
|
}
|
|
}
|
|
|
|
fn configure_handler() -> BotHandler {
|
|
match env::var("OPENAI_TOKEN") {
|
|
Ok(token) => {
|
|
log::debug!("OpenAI functionality enabled");
|
|
let default_model = env::var("OPENAI_MODEL").unwrap_or_else(|_| "gpt-4o-mini".to_string());
|
|
let base_url = env::var("OPENAI_BASE_URL").unwrap();
|
|
BotHandler {
|
|
oai: Some(OAI {
|
|
client: reqwest::Client::new(),
|
|
base_url,
|
|
token,
|
|
max_conversation_history: 30,
|
|
max_tokens: 8192,
|
|
default_model,
|
|
}),
|
|
}
|
|
}
|
|
Err(_) => {
|
|
log::warn!("OpenAI functionality disabled");
|
|
BotHandler { oai: None }
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn signal_shutdown(shard_manager: Arc<ShardManager>) {
|
|
tokio::signal::ctrl_c()
|
|
.await
|
|
.expect("Failed to listen for shutdown signal");
|
|
shard_manager.shutdown_all().await;
|
|
log::info!("Bot shutdown gracefully.");
|
|
}
|