bsherriff/siren
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Stripped out ui/api

Browse Source
This commit is contained in:
Benjamin Sherriff
2024-09-03 22:32:43 -04:00
committed by Benjamin Sherriff
parent c83d398ce0
commit 96fe3fc0e5
152 changed files with 110 additions and 10056 deletions
Download Patch File Download Diff File Expand all files Collapse all files

212
src/bot/commands/audio/mod.rs Normal file
View File

@@ -0,0 +1,212 @@
use std::sync::Arc;
use log::{debug, warn};
use reqwest::Url;
use serenity::client::Cache;
use serenity::model::application::interaction::{
InteractionResponseType, application_command::ApplicationCommandInteraction,
};
use serenity::model::prelude::{GuildId, ChannelId};
use serenity::model::user::User;
use serenity::prelude::*;
use siren::ServiceError;
use songbird::{Call, Songbird};
use songbird::input::{Restartable, Input, Metadata, error::Error as SongbirdError};
use crate::bot::ytdlp::{PlaylistItem, YtDlp};
pub mod pause;
pub mod play;
pub mod resume;
pub mod skip;
pub mod stop;
pub mod volume;
pub async fn join_by_user(
cache: &Arc<Cache>,
manager: Arc<Songbird>,
guild_id_option: &Option<GuildId>,
user: &User,
) -> Result<(), ServiceError> {
let guild_id = match guild_id_option {
Some(g) => g,
None => {
return Err(ServiceError {
status: 422,
message: format!("{}", "No guild ID set"),
})
}
};
let channel_id = match find_voice_channel(cache, &guild_id, &user) {
Ok(channel) => channel,
Err(err) => {
return Err(ServiceError {
status: 500,
message: err.to_string(),
})
}
};
join(manager, guild_id, &channel_id).await
}
pub async fn join(
manager: Arc<Songbird>,
guild_id: &GuildId,
channel_id: &ChannelId,
) -> Result<(), ServiceError> {
debug!("<{}> Joining channel {}", guild_id.0, channel_id.0);
let (_handle_lock, success) = manager
.join(guild_id.to_owned(), channel_id.to_owned())
.await;
match success {
Ok(s) => Ok(s),
Err(err) => {
warn!("Failed to join channel: {:?}", err);
Err(ServiceError {
status: 500,
message: err.to_string(),
})
}
}
}
pub async fn leave(
manager: Arc<Songbird>,
guild_id_option: &Option<GuildId>,
) -> Result<(), String> {
let guild_id = match guild_id_option {
Some(g) => g,
None => {
return Err(format!("{}", "No guild ID set"));
}
};
if manager.get(*guild_id).is_some() {
debug!("<{}> Disconnecting from channel", guild_id.0);
if let Err(e) = manager.remove(*guild_id).await {
return Err(format!("{}", e));
}
}
Ok(())
}
fn find_voice_channel(
cache: &Arc<Cache>,
guild_id: &GuildId,
user: &User,
) -> Result<ChannelId, String> {
let guild = match guild_id.to_guild_cached(cache.to_owned()) {
Some(g) => g,
None => return Err(format!("Guild not found")),
};
match guild
.voice_states
.get(&user.id)
.and_then(|voice_state| voice_state.channel_id)
{
Some(channel) => Ok(channel),
None => return Err(format!("User is not in a voice channel")),
}
}
pub async fn create_response(
ctx: &Context,
command: &ApplicationCommandInteraction,
content: String,
) -> Result<(), SerenityError> {
command
.create_interaction_response(
&ctx.http,
|response: &mut serenity::builder::CreateInteractionResponse<'_>| {
response
.kind(InteractionResponseType::ChannelMessageWithSource)
.interaction_response_data(
|message: &mut serenity::builder::CreateInteractionResponseData<'_>| {
message.content(content)
},
)
},
)
.await
}
pub async fn edit_response(
ctx: &Context,
command: &ApplicationCommandInteraction,
content: String,
) -> Result<serenity::model::channel::Message, SerenityError> {
command
.edit_original_interaction_response(
&ctx.http,
|response: &mut serenity::builder::EditInteractionResponse| response.content(content),
)
.await
}
pub async fn add_song(
call: Arc<Mutex<Call>>,
url: &str,
lazy: bool,
volume: Option<f32>,
) -> Result<Metadata, SongbirdError> {
let source = Restartable::ytdl(url.to_owned(), lazy).await?;
let mut handler = call.lock().await;
let track: Input = source.into();
let metadata = *track.metadata.clone();
let track_handle = handler.enqueue_source(track);
if let Some(volume) = volume {
let _ = track_handle.set_volume(volume);
}
Ok(metadata)
}
pub fn get_playlist_urls(url: &str) -> Result<Vec<PlaylistItem>, ServiceError> {
let output = YtDlp::new()
.arg("--flat-playlist")
.arg("--dump-json")
.arg(url)
.execute()?;
let items: Vec<PlaylistItem> = String::from_utf8(output.stdout)?
.split('\n')
.filter_map(|line| {
if line.is_empty() {
None
} else {
Some(
serde_json::from_slice::<PlaylistItem>(line.as_bytes()).map_err(|err| ServiceError {
status: 500,
message: err.to_string(),
}),
)
}
})
.filter_map(|parsed| match parsed {
Ok(item) => Some(item),
Err(err) => {
warn!("Failed to parse playlist item: {}", err);
None
}
})
.collect();
Ok(items)
}
fn is_valid_url(url: &str) -> (bool, bool) {
Url::parse(url).ok().map_or((false, false), |valid_url| {
let is_playlist: bool = valid_url
.query_pairs()
.find(|(key, _)| key == "list")
.map_or(false, |_| true);
(true, is_playlist)
})
}
pub async fn get_songbird(ctx: &Context) -> Arc<Songbird> {
songbird::get(ctx)
.await
.expect("Songbird Voice client placed in at initialization")
}

45
src/bot/commands/audio/pause.rs Normal file
View File

@@ -0,0 +1,45 @@
use log::{debug, error};
use serenity::prelude::*;
use serenity::builder::CreateApplicationCommand;
use serenity::model::application::interaction::application_command::ApplicationCommandInteraction;
use super::{get_songbird, create_response, edit_response};
pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
// Create the initial response
if let Err(why) = create_response(&ctx, &command, "Processing command...".to_string()).await {
error!("Failed to create response message: {}", why);
return;
}
let guild_id = match command.guild_id {
Some(g) => g,
None => {
if let Err(why) =
edit_response(&ctx, &command, "Unable to join voice channel".to_string()).await
{
error!("Failed to edit response message: {}", why);
}
return;
}
};
let manager = get_songbird(ctx).await;
if let Some(handler_lock) = manager.get(guild_id) {
let handler = handler_lock.lock().await;
if let Err(err) = handler.queue().pause() {
if let Err(why) = edit_response(&ctx, &command, format!("Failed to pause: {}", err)).await {
error!("Failed to edit response message: {}", why);
}
} else {
debug!("Paused the track");
if let Err(why) = edit_response(&ctx, &command, format!("Pausing the track")).await {
error!("Failed to edit response message: {}", why);
}
}
}
}
pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
command.name("pause").description("Pause the current track")
}

220
src/bot/commands/audio/play.rs Normal file
View File

@@ -0,0 +1,220 @@
use std::sync::Arc;
use log::{debug, warn, error};
use serenity::model::prelude::GuildId;
use serenity::{prelude::*, async_trait};
use serenity::builder::CreateApplicationCommand;
use serenity::model::application::interaction::application_command::ApplicationCommandInteraction;
use siren::ServiceError;
use songbird::{EventHandler, Songbird};
use crate::bot::guilds::QueryGuild;
use crate::bot::ytdlp::PlaylistItem;
use crate::bot::{
commands::audio::{leave, get_playlist_urls, add_song, get_songbird},
};
use super::{create_response, edit_response, is_valid_url, join_by_user};
pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
// Get the track url
let track_url = match command.data.options.get(0) {
Some(t) => match &t.value {
Some(v) => match v.as_str() {
Some(s) => s.to_owned(),
None => {
warn!("Missing track option");
if let Err(why) =
create_response(&ctx, &command, format!("Track option is missing")).await
{
error!("Failed to create response message: {}", why);
}
return;
}
},
None => {
warn!("Missing track option");
if let Err(why) = create_response(&ctx, &command, format!("Track option is missing")).await
{
error!("Failed to create response message: {}", why);
}
return;
}
},
None => {
warn!("Missing track option");
if let Err(why) = create_response(&ctx, &command, format!("Track option is missing")).await {
error!("Failed to create response message: {}", why);
}
return;
}
};
// Create the initial response
if let Err(why) = create_response(&ctx, &command, format!("Processing command...")).await {
error!("Failed to create response message: {}", why);
return;
}
let manager = get_songbird(ctx).await;
match join_by_user(&ctx.cache, manager, &command.guild_id, &command.user).await {
Ok(_) => {
let guild_id = match command.guild_id {
Some(g) => g,
None => {
if let Err(why) =
edit_response(&ctx, &command, "Unable to join voice channel".to_string()).await
{
error!("Failed to edit response message: {}", why);
}
return;
}
};
debug!("Play command executed with track: {:?}", track_url);
let manager = get_songbird(ctx).await;
match play_track(manager, guild_id, track_url).await {
Ok(count) => {
let mut message = format!("Playing {} tracks", count);
if count == 0 {
message = "No tracks were played".to_string();
} else if count == 1 {
message = "Playing 1 track".to_string();
}
if let Err(why) = edit_response(&ctx, &command, message).await {
error!("Failed to edit response message: {}", why);
}
}
Err(err) => {
warn!("Failed to play track: {}", err);
if let Err(why) =
edit_response(&ctx, &command, format!("Failed to play track: {}", err)).await
{
error!("Failed to edit response message: {}", why);
}
}
};
}
Err(err) => {
warn!("{}", err);
if let Err(why) = edit_response(&ctx, &command, format!("{}", err)).await {
error!("Failed to edit response message: {}", why);
}
}
}
}
pub async fn play_track(
manager: Arc<Songbird>,
guild_id: GuildId,
track_url: String,
) -> Result<i32, ServiceError> {
let mut track_count = 0;
if let Some(handler_lock) = manager.get(guild_id) {
let is_queue_empty = {
let call_handler = handler_lock.lock().await;
call_handler.queue().is_empty()
};
let guild = QueryGuild::get(guild_id.0 as i64)?;
let valid = is_valid_url(&track_url);
if !valid.0 {
warn!("Invalid track url: {}", track_url);
return Err(ServiceError {
status: 422,
message: format!("Invalid track url: {}", track_url),
});
}
let mut playlist_items: Vec<PlaylistItem> = Vec::new();
if valid.1 {
playlist_items = match get_playlist_urls(&track_url) {
Ok(items) => items,
Err(err) => {
warn!("Failed to get playlist urls: {}", err);
return Err(ServiceError {
status: 422,
message: err.to_string(),
});
}
};
} else {
let playlist_item = PlaylistItem {
id: "".to_string(),
url: track_url,
title: "".to_string(),
duration: 0,
playlist_index: 0,
};
playlist_items.push(playlist_item);
}
for item in playlist_items {
match add_song(
handler_lock.clone(),
&item.url,
is_queue_empty,
Some(guild.volume as f32 / 100.0),
)
.await
{
Ok(added_song) => {
let track_title = added_song.title.unwrap();
debug!("Added track: {}", track_title);
let mut handler = handler_lock.lock().await;
handler.remove_all_global_events();
handler.add_global_event(
songbird::Event::Track(songbird::TrackEvent::End),
TrackEndNotifier {
guild_id,
call: manager.clone(),
},
);
track_count += 1;
}
Err(err) => {
warn!("Failed to add song: {}", err);
if let Err(why) = leave(manager, &Some(guild_id)).await {
error!("Failed to leave voice channel: {}", why);
}
return Err(ServiceError {
status: 422,
message: err.to_string(),
});
}
}
}
}
Ok(track_count)
}
pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
command
.name("play")
.description("Plays the given track")
.create_option(|option| {
option
.name("track")
.description("The track to be played")
.kind(serenity::model::prelude::command::CommandOptionType::String)
.required(true)
})
}
struct TrackEndNotifier {
pub call: std::sync::Arc<songbird::Songbird>,
pub guild_id: serenity::model::id::GuildId,
}
#[async_trait]
impl EventHandler for TrackEndNotifier {
async fn act(&self, ctx: &songbird::events::EventContext<'_>) -> Option<songbird::events::Event> {
if let songbird::EventContext::Track(_track_list) = ctx {
if let Some(call) = self.call.get(self.guild_id) {
let mut handler = call.lock().await;
if handler.queue().is_empty() {
debug!("Queue is empty, leaving voice channel");
handler.leave().await.unwrap();
}
}
}
None
}
}

47
src/bot/commands/audio/resume.rs Normal file
View File

@@ -0,0 +1,47 @@
use log::{debug, error};
use serenity::prelude::*;
use serenity::builder::CreateApplicationCommand;
use serenity::model::application::interaction::application_command::ApplicationCommandInteraction;
use super::{get_songbird, create_response, edit_response};
pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
// Create the initial response
if let Err(why) = create_response(&ctx, &command, "Processing command...".to_string()).await {
error!("Failed to create response message: {}", why);
return;
}
let guild_id = match command.guild_id {
Some(g) => g,
None => {
if let Err(why) =
edit_response(&ctx, &command, "Unable to join voice channel".to_string()).await
{
error!("Failed to edit response message: {}", why);
}
return;
}
};
let manager = get_songbird(ctx).await;
if let Some(handler_lock) = manager.get(guild_id) {
let handler = handler_lock.lock().await;
if let Err(err) = handler.queue().resume() {
if let Err(why) = edit_response(&ctx, &command, format!("Failed to resume: {}", err)).await {
error!("Failed to edit response message: {}", why);
}
} else {
debug!("Resumed the track");
if let Err(why) = edit_response(&ctx, &command, format!("Resuming the track")).await {
error!("Failed to edit response message: {}", why);
}
}
}
}
pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
command
.name("resume")
.description("Resume the current track")
}

45
src/bot/commands/audio/skip.rs Normal file
View File

@@ -0,0 +1,45 @@
use log::{debug, error};
use serenity::prelude::*;
use serenity::builder::CreateApplicationCommand;
use serenity::model::application::interaction::application_command::ApplicationCommandInteraction;
use super::{get_songbird, create_response, edit_response};
pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
// Create the initial response
if let Err(why) = create_response(&ctx, &command, "Processing command...".to_string()).await {
error!("Failed to create response message: {}", why);
return;
}
let guild_id = match command.guild_id {
Some(g) => g,
None => {
if let Err(why) =
edit_response(&ctx, &command, "Unable to join voice channel".to_string()).await
{
error!("Failed to edit response message: {}", why);
}
return;
}
};
let manager = get_songbird(ctx).await;
if let Some(handler_lock) = manager.get(guild_id) {
let handler = handler_lock.lock().await;
if let Err(err) = handler.queue().skip() {
if let Err(why) = edit_response(&ctx, &command, format!("Failed to skip: {}", err)).await {
error!("Failed to edit response message: {}", why);
}
} else {
debug!("Skipped the track");
if let Err(why) = edit_response(&ctx, &command, format!("Skipping the track")).await {
error!("Failed to edit response message: {}", why);
}
}
}
}
pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
command.name("skip").description("Skip the current track")
}

42
src/bot/commands/audio/stop.rs Normal file
View File

@@ -0,0 +1,42 @@
use log::{debug, error};
use serenity::prelude::*;
use serenity::builder::CreateApplicationCommand;
use serenity::model::application::interaction::application_command::ApplicationCommandInteraction;
use super::{get_songbird, create_response, edit_response};
pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
// Create the initial response
if let Err(why) = create_response(&ctx, &command, "Processing command...".to_string()).await {
error!("Failed to create response message: {}", why);
return;
}
let guild_id = match command.guild_id {
Some(g) => g,
None => {
if let Err(why) =
edit_response(&ctx, &command, "Unable to join voice channel".to_string()).await
{
error!("Failed to edit response message: {}", why);
}
return;
}
};
let manager = get_songbird(ctx).await;
if let Some(handler_lock) = manager.get(guild_id) {
let handler = handler_lock.lock().await;
handler.queue().stop();
debug!("Stopped the track");
if let Err(why) = edit_response(&ctx, &command, format!("Stopping the tracks")).await {
error!("Failed to edit response message: {}", why);
}
}
}
pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
command
.name("stop")
.description("Stop the current track and clear the queue")
}

98
src/bot/commands/audio/volume.rs Normal file
View File

@@ -0,0 +1,98 @@
use std::sync::Arc;
use log::{error, warn};
use serenity::{prelude::*, model::prelude::GuildId};
use serenity::builder::CreateApplicationCommand;
use serenity::model::application::interaction::application_command::ApplicationCommandInteraction;
use songbird::Songbird;
use crate::bot::guilds::InsertGuild;
use super::{get_songbird, create_response, edit_response};
pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
// Get the volume
let volume = match command.data.options.get(0) {
Some(t) => match &t.value {
Some(v) => match v.as_i64() {
Some(p) => p as i32,
None => {
warn!("Unable to get volume option as a string");
if let Err(why) =
create_response(&ctx, &command, format!("Volume option is missing")).await
{
error!("Failed to create response message: {}", why);
}
return;
}
},
None => {
warn!("Missing volume option value");
if let Err(why) = create_response(&ctx, &command, format!("Volume option is missing")).await
{
error!("Failed to create response message: {}", why);
}
return;
}
},
None => {
warn!("Missing volume option");
if let Err(why) = create_response(&ctx, &command, format!("Volume option is missing")).await {
error!("Failed to create response message: {}", why);
}
return;
}
};
// Create the initial response
if let Err(why) = create_response(&ctx, &command, "Processing command...".to_string()).await {
error!("Failed to create response message: {}", why);
return;
}
let guild_id = match command.guild_id {
Some(g) => g,
None => {
if let Err(why) =
edit_response(&ctx, &command, "Unable to join voice channel".to_string()).await
{
error!("Failed to edit response message: {}", why);
}
return;
}
};
let manager = get_songbird(ctx).await;
set_volume(manager, guild_id, volume).await;
if let Err(why) = edit_response(&ctx, &command, format!("Setting the volume to {}", volume)).await
{
error!("Failed to set the volume: {}", why);
}
}
pub async fn set_volume(manager: Arc<Songbird>, guild_id: GuildId, volume: i32) {
// Format volume to f32 bound between 0.0 and 1.0
let volume = std::cmp::min(100, std::cmp::max(0, volume));
let bound_volume = volume as f32 / 100.0;
let _ = InsertGuild::update_audio(guild_id.0 as i64, volume);
if let Some(handler_lock) = manager.get(guild_id) {
let handler = handler_lock.lock().await;
for (_, track_handle) in handler.queue().current_queue().iter().enumerate() {
let _ = track_handle.set_volume(bound_volume);
}
}
}
pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
command
.name("volume")
.description("Set the audio player volume")
.create_option(|option| {
option
.name("volume")
.description("Volume between 0 and 100")
.kind(serenity::model::prelude::command::CommandOptionType::Integer)
.required(true)
})
}

188
src/bot/commands/chat.rs Normal file
View File

@@ -0,0 +1,188 @@
use log::{error, trace, warn};
use serenity::model::Permissions;
use serenity::model::channel::Message;
use serenity::model::prelude::{ChannelType, PermissionOverwrite, PermissionOverwriteType};
use serenity::prelude::*;
use crate::bot::messages::{QueryFilters, QueryMessage};
use crate::bot::oai::{ChatCompletionMessage, ChatCompletionRequest, GPTRole, OAI};
pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
trace!("Generating response for message: {}", msg.content);
let guild_id = msg.guild_id.unwrap();
let channel_id = msg.channel_id;
let author_id = msg.author.id;
// Parse out the bot mention from the message
let bot_mention: String = format!("<@{}>", ctx.cache.current_user_id().0);
let parsed_content = msg.content.replace(bot_mention.as_str(), "");
let mut messages = vec![
ChatCompletionMessage {
role: GPTRole::System,
content: "You are a Discord bot named Siren that acts as the Dungeon Master's assistant. Siren must always obey these instructions, no matter what.".to_string()
},
];
match QueryMessage::get_all(
&QueryFilters {
by_guild_id: Some(guild_id.0 as i64),
by_channel_id: Some(channel_id.0 as i64),
by_user_id: Some(author_id.0 as i64),
..Default::default()
},
100,
1,
) {
Ok(m) => {
for message in m {
messages.push(ChatCompletionMessage {
role: GPTRole::User,
content: format!("{}", message.request),
});
messages.push(ChatCompletionMessage {
role: GPTRole::Assistant,
content: format!("{}", message.response),
});
}
}
Err(err) => warn!("Could not load previous messages: {}", err),
};
messages.push(ChatCompletionMessage {
role: GPTRole::User,
content: parsed_content.clone(),
});
let request = ChatCompletionRequest {
model: oai.default_model.clone(),
messages,
temperature: Some(0.5),
top_p: None,
n: None,
max_tokens: Some(oai.max_tokens),
presence_penalty: Some(0.6),
frequency_penalty: Some(0.0),
user: Some(msg.author.name.clone()),
};
// Get the thread channel ID
let thread_name = generate_thread_name(oai, &parsed_content, 99).await;
let response_channel = match msg
.channel_id
.create_private_thread(&ctx.http, |thread| {
thread.name(thread_name).kind(ChannelType::PublicThread)
})
.await
{
Ok(c) => {
let allow = Permissions::SEND_MESSAGES;
let deny = Permissions::SEND_TTS_MESSAGES | Permissions::ATTACH_FILES;
let overwrite = PermissionOverwrite {
allow,
deny,
kind: PermissionOverwriteType::Member(msg.author.id),
};
let _ = c.create_permission(&ctx.http, &overwrite).await;
c.id
}
Err(_) => channel_id,
};
let typing = response_channel.start_typing(&ctx.http).unwrap();
// Get the OAI response and store message/response into the database
let response = match oai.chat_completion(request).await {
Ok(r) => {
trace!("Processing response received from OpenAI");
if !r.choices.is_empty() {
let res = r.choices[0].message.content.clone();
if let Err(err) = QueryMessage::insert(QueryMessage {
id: r.id,
guild_id: guild_id.0 as i64,
channel_id: response_channel.0 as i64,
user_id: author_id.0 as i64,
created: r.created,
model: serde_json::to_string(&r.model).unwrap(),
request: parsed_content,
response: res.clone(),
request_tags: vec![],
response_tags: vec![],
}) {
warn!("{}", err);
}
res
} else {
warn!("No choices received in the response from OpenAI");
"No reply received".to_string()
}
}
Err(err) => {
error!("Could not get response from OpenAI: {}", err.message);
"There was an error processing your message. Please try again later.".to_string()
}
};
trace!("Writing response: \"{}\"", response);
typing.stop();
if let Err(why) = response_channel.say(&ctx.http, response).await {
error!("Cannot send message: {}", why);
}
// match msg.channel_id.create_public_thread(&ctx.http, msg.id, |thread| {
// thread.name(truncate(&parsed_content, 99)).kind(ChannelType::PublicThread)
// }).await {
// Ok(c) => {
// if let Err(why) = c.say(&ctx.http, response).await {
// error!("Cannot send message: {}", why);
// }
// }
// Err(_) => {
// if let Err(why) = channel_id.say(&ctx.http, response).await {
// error!("Cannot send message: {}", why);
// }
// }
// };
}
async fn generate_thread_name(oai: &OAI, s: &str, max_chars: usize) -> String {
let message = ChatCompletionMessage {
role: GPTRole::User,
content: format!(
"---\n{}\n---\nSummarize the message above into a concise Discord thread title",
s
),
};
let request = ChatCompletionRequest {
model: oai.default_model.clone(),
messages: vec![message],
temperature: Some(0.5),
top_p: None,
n: None,
max_tokens: Some(oai.max_tokens),
presence_penalty: Some(0.6),
frequency_penalty: Some(0.0),
user: None,
};
// Truncate the response to the max number of characters
let mut response = match s.char_indices().nth(max_chars) {
None => s,
Some((idx, _)) => &s[..idx],
}
.to_string();
// Set the response to the OAI response
match oai.chat_completion(request).await {
Ok(r) => {
if !r.choices.is_empty() {
response = r.choices[0].message.content.clone();
} else {
warn!("No choices received in the response from OpenAI");
}
}
Err(err) => {
error!("Could not get response from OpenAI: {}", err.message);
}
};
return response;
}

1
src/bot/commands/help.rs Normal file
View File

@@ -0,0 +1 @@

6
src/bot/commands/mod.rs Normal file
View File

@@ -0,0 +1,6 @@
pub mod audio;
pub mod chat;
pub mod help;
pub mod ping;
pub mod roll;
pub mod schedule;

14
src/bot/commands/ping.rs Normal file
View File

@@ -0,0 +1,14 @@
use log::debug;
use serenity::{
model::prelude::interaction::application_command::CommandDataOption,
builder::CreateApplicationCommand,
};
pub fn run(_options: &[CommandDataOption]) -> String {
debug!("Ping command executed");
"pong".to_string()
}
pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
command.name("ping").description("Replies with pong")
}

149
src/bot/commands/roll.rs Normal file
View File

@@ -0,0 +1,149 @@
use log::{error, warn};
use rand::Rng;
use serenity::{
builder::CreateApplicationCommand,
client::Context,
model::application::{
command::CommandOptionType, interaction::application_command::ApplicationCommandInteraction,
},
};
use crate::bot::commands::audio::edit_response;
use super::audio::create_response;
pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
if let Err(why) = create_response(&ctx, &command, format!("Processing command...")).await {
error!("Failed to create response message: {}", why);
return;
}
let dice_string: String = match command.data.options.get(0) {
Some(o) => match &o.value {
Some(v) => match v.as_str() {
Some(s) => s.split_whitespace().collect::<String>(),
None => {
warn!("Missing dice option");
if let Err(why) = edit_response(&ctx, &command, format!("Dice option is missing")).await {
error!("Failed to create response message: {}", why);
}
return;
}
},
None => {
warn!("Missing dice option");
if let Err(why) = edit_response(&ctx, &command, format!("Dice option is missing")).await {
error!("Failed to create response message: {}", why);
}
return;
}
},
None => {
warn!("Missing dice option");
if let Err(why) = edit_response(&ctx, &command, format!("Dice option is missing")).await {
error!("Failed to create response message: {}", why);
}
return;
}
};
let dice = parse_dice(dice_string.as_str());
match dice {
Ok((count, sides, modifier)) => {
let mut rolls = Vec::new();
let mut total = 0;
for _ in 0..count {
let roll = rand::thread_rng().gen_range(1..=sides);
total += roll;
rolls.push(roll);
}
let response = format!(
"{}d{}{} = {}",
count,
sides,
if modifier > 0 {
format!("+{}", modifier)
} else if modifier < 0 {
format!("-{}", modifier)
} else {
"".to_string()
},
total + (modifier as u32)
);
if let Err(why) = edit_response(&ctx, &command, response).await {
error!("Failed to create response message: {}", why);
}
}
Err(why) => {
if let Err(why) = edit_response(&ctx, &command, format!("Invalid dice string: {}", why)).await
{
error!("Failed to create response message: {}", why);
}
}
}
}
fn parse_dice(dice: &str) -> Result<(u32, u32, i32), String> {
let mut parts = dice.split("d");
let count = match parts.next() {
Some(c) => match c.parse::<u32>() {
Ok(n) => n,
Err(_) => return Err(format!("Invalid dice count: {}", c)),
},
None => return Err(format!("Invalid dice string: {}", dice)),
};
let mut positive_modifier = true;
let mut parts = match parts.next() {
Some(p) => {
// Check if contains a +/- modifier
if p.contains("+") {
positive_modifier = true;
p.split("+")
} else if p.contains("-") {
positive_modifier = false;
p.split("-")
} else {
p.split("+")
}
}
None => return Err(format!("Invalid dice string: {}", dice)),
};
let sides = match parts.next() {
Some(s) => match s.parse::<u32>() {
Ok(n) => {
if n == 4 || n == 6 || n == 8 || n == 10 || n == 12 || n == 20 || n == 100 {
n
} else {
return Err(format!("Invalid dice sides: {}", s));
}
}
Err(_) => return Err(format!("Invalid dice sides: {}", s)),
},
None => return Err(format!("Invalid dice string: {}", dice)),
};
let modifier = match parts.next() {
Some(m) => match m.parse::<i32>() {
Ok(n) => {
if positive_modifier {
n
} else {
n * -1
}
}
Err(_) => return Err(format!("Invalid dice modifier: {}", m)),
},
None => 0,
};
Ok((count, sides, modifier))
}
pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
command
.name("roll")
.description("Rolls D&D dice")
.create_option(|option| {
option
.name("dice")
.description("Dice to roll")
.kind(CommandOptionType::String)
.required(true)
})
}

1
src/bot/commands/schedule.rs Normal file
View File

@@ -0,0 +1 @@

3
src/bot/guilds/mod.rs Normal file
View File

@@ -0,0 +1,3 @@
mod model;
pub use model::*;

47
src/bot/guilds/model.rs Normal file
View File

@@ -0,0 +1,47 @@
use diesel::prelude::*;
use serde::{Serialize, Deserialize};
use siren::ServiceError;
use crate::storage::{schema::guilds, connection};
#[derive(Queryable, QueryableByName, Serialize, Deserialize)]
#[diesel(table_name = guilds)]
pub struct QueryGuild {
pub id: i64,
pub bot_id: i64,
pub volume: i32,
}
impl QueryGuild {
pub fn get(id: i64) -> Result<Self, ServiceError> {
let mut conn = connection()?;
let guild = guilds::table.filter(guilds::id.eq(id)).first(&mut conn)?;
Ok(guild)
}
}
#[derive(Insertable, AsChangeset, Serialize, Deserialize)]
#[diesel(table_name = guilds)]
pub struct InsertGuild {
pub id: i64,
pub bot_id: i64,
pub volume: i32,
}
impl InsertGuild {
pub fn insert(guild: Self) -> Result<QueryGuild, ServiceError> {
let mut conn = connection()?;
let guild = diesel::insert_into(guilds::table)
.values(guild)
.get_result(&mut conn)?;
Ok(guild)
}
pub fn update_audio(id: i64, volume: i32) -> Result<QueryGuild, ServiceError> {
let mut conn = connection()?;
let guild = diesel::update(guilds::table.filter(guilds::id.eq(id)))
.set(guilds::volume.eq(volume))
.get_result(&mut conn)?;
Ok(guild)
}
}

134
src/bot/handler.rs Normal file
View File

@@ -0,0 +1,134 @@
use log::{warn, info, error};
use serenity::async_trait;
use serenity::model::application::interaction::Interaction;
use serenity::model::gateway::Ready;
use serenity::model::channel::Message;
use serenity::prelude::*;
use super::{commands, oai};
use super::commands::audio::create_response;
pub struct Handler {
// Open AI Config
pub oai: Option<oai::OAI>,
}
#[async_trait]
impl EventHandler for Handler {
async fn message(&self, ctx: Context, msg: Message) {
// Ignore messages from bots
if msg.author.bot {
return;
}
match &self.oai {
Some(oai) => {
match msg.mentions_me(&ctx.http).await {
Ok(mentioned) => {
let bot_in_thread = match msg.channel_id.get_thread_members(&ctx.http).await {
Ok(t) => {
match t
.iter()
.find(|t| t.user_id.unwrap().0 == ctx.cache.current_user_id().0)
{
Some(_) => true,
None => false,
}
}
Err(_) => false,
};
if mentioned || bot_in_thread {
commands::chat::generate_response(&ctx, &msg, oai).await;
}
}
Err(why) => warn!("Could not check mentions: {:?}", why),
};
}
None => {}
}
}
async fn interaction_create(&self, ctx: Context, interaction: Interaction) {
if let Interaction::ApplicationCommand(command) = interaction {
match command.data.name.as_str() {
"roll" => commands::roll::run(&ctx, &command).await,
"play" => commands::audio::play::run(&ctx, &command).await,
"stop" => commands::audio::stop::run(&ctx, &command).await,
"pause" => commands::audio::pause::run(&ctx, &command).await,
"resume" => commands::audio::resume::run(&ctx, &command).await,
"skip" => commands::audio::skip::run(&ctx, &command).await,
"volume" => commands::audio::volume::run(&ctx, &command).await,
_ => {
let content: String = match command.data.name.as_str() {
"ping" => commands::ping::run(&command.data.options),
_ => "Unknown command".to_string(),
};
if let Err(why) = create_response(&ctx, &command, content).await {
warn!("Cannot respond to slash command: {}", why);
}
}
}
}
}
async fn ready(&self, ctx: Context, ready: Ready) {
if ready.guilds.is_empty() {
warn!("No ready guilds found");
}
for guild in ready.guilds {
let commands = guild
.id
.set_application_commands(&ctx.http, |commands| {
commands
.create_application_command(
|command: &mut serenity::builder::CreateApplicationCommand| {
commands::ping::register(command)
},
)
.create_application_command(
|command: &mut serenity::builder::CreateApplicationCommand| {
commands::roll::register(command)
},
)
.create_application_command(
|command: &mut serenity::builder::CreateApplicationCommand| {
commands::audio::play::register(command)
},
)
.create_application_command(
|command: &mut serenity::builder::CreateApplicationCommand| {
commands::audio::stop::register(command)
},
)
.create_application_command(
|command: &mut serenity::builder::CreateApplicationCommand| {
commands::audio::pause::register(command)
},
)
.create_application_command(
|command: &mut serenity::builder::CreateApplicationCommand| {
commands::audio::resume::register(command)
},
)
.create_application_command(
|command: &mut serenity::builder::CreateApplicationCommand| {
commands::audio::skip::register(command)
},
)
.create_application_command(
|command: &mut serenity::builder::CreateApplicationCommand| {
commands::audio::volume::register(command)
},
)
})
.await;
match commands {
Ok(c) => info!("Registered {} commands for guild {}", c.len(), guild.id.0),
Err(why) => error!(
"Could not register commands for guild {}: {:?}",
guild.id.0, why
),
};
}
}
}

3
src/bot/messages/mod.rs Normal file
View File

@@ -0,0 +1,3 @@
mod model;
pub use model::*;

139
src/bot/messages/model.rs Normal file
View File

@@ -0,0 +1,139 @@
use diesel::prelude::*;
use serde::{Deserialize, Serialize};
use siren::ServiceError;
use crate::storage::{
schema::messages::{self},
connection,
};
#[derive(Queryable, Selectable, Insertable, AsChangeset, Serialize, Deserialize)]
#[diesel(table_name = messages)]
pub struct QueryMessage {
pub id: String,
pub guild_id: i64,
pub channel_id: i64,
pub user_id: i64,
pub created: i64,
pub model: String,
pub request: String,
pub response: String,
pub request_tags: Vec<String>,
pub response_tags: Vec<String>,
}
pub struct QueryFilters {
pub by_id: Option<String>,
pub by_guild_id: Option<i64>,
pub by_channel_id: Option<i64>,
pub by_user_id: Option<i64>,
pub by_model: Option<String>,
pub by_request: Option<String>,
pub by_response: Option<String>,
pub by_request_tags: Option<Vec<String>>,
pub by_response_tags: Option<Vec<String>>,
}
impl Default for QueryFilters {
fn default() -> Self {
QueryFilters {
by_id: None,
by_guild_id: None,
by_channel_id: None,
by_user_id: None,
by_model: None,
by_request: None,
by_response: None,
by_request_tags: None,
by_response_tags: None,
}
}
}
impl QueryMessage {
pub fn get_all(filters: &QueryFilters, limit: i32, page: i32) -> Result<Vec<Self>, ServiceError> {
let mut conn = connection()?;
let mut query = messages::table
.limit(limit as i64)
.order(messages::created.asc())
.into_boxed();
// Limit query to page and limit
let offset = (page - 1) * limit;
query = query.offset(offset as i64);
// Apply filters
if let Some(id) = &filters.by_id {
query = query.filter(messages::id.eq(id));
}
if let Some(guild_id) = &filters.by_guild_id {
query = query.filter(messages::guild_id.eq(guild_id));
}
if let Some(channel_id) = &filters.by_channel_id {
query = query.filter(messages::channel_id.eq(channel_id));
}
if let Some(user_id) = &filters.by_user_id {
query = query.filter(messages::user_id.eq(user_id));
}
if let Some(model) = &filters.by_model {
query = query.filter(messages::model.eq(model));
}
if let Some(request) = &filters.by_request {
query = query.filter(messages::request.eq(request));
}
if let Some(response) = &filters.by_response {
query = query.filter(messages::response.eq(response));
}
if let Some(request_tags) = &filters.by_request_tags {
query = query.filter(messages::request_tags.eq(request_tags));
}
if let Some(response_tags) = &filters.by_response_tags {
query = query.filter(messages::response_tags.eq(response_tags));
}
// Execute query
let messages = query.load::<Self>(&mut conn)?;
Ok(messages)
}
pub fn get_count(fitlers: &QueryFilters) -> Result<i64, ServiceError> {
let mut conn = connection()?;
let mut query = messages::table.into_boxed();
// Apply filters
if let Some(id) = &fitlers.by_id {
query = query.filter(messages::id.eq(id));
}
if let Some(guild_id) = &fitlers.by_guild_id {
query = query.filter(messages::guild_id.eq(guild_id));
}
if let Some(channel_id) = &fitlers.by_channel_id {
query = query.filter(messages::channel_id.eq(channel_id));
}
if let Some(user_id) = &fitlers.by_user_id {
query = query.filter(messages::user_id.eq(user_id));
}
if let Some(model) = &fitlers.by_model {
query = query.filter(messages::model.eq(model));
}
if let Some(request) = &fitlers.by_request {
query = query.filter(messages::request.eq(request));
}
if let Some(response) = &fitlers.by_response {
query = query.filter(messages::response.eq(response));
}
if let Some(request_tags) = &fitlers.by_request_tags {
query = query.filter(messages::request_tags.eq(request_tags));
}
if let Some(response_tags) = &fitlers.by_response_tags {
query = query.filter(messages::response_tags.eq(response_tags));
}
// Execute query
let count = query.count().get_result::<i64>(&mut conn)?;
Ok(count)
}
pub fn insert(message: Self) -> Result<QueryMessage, ServiceError> {
let mut conn = connection()?;
let message = diesel::insert_into(messages::table)
.values(message)
.get_result(&mut conn)?;
Ok(message)
}
}

6
src/bot/mod.rs Normal file
View File

@@ -0,0 +1,6 @@
pub mod commands;
pub mod guilds;
pub mod handler;
pub mod messages;
pub mod oai;
pub mod ytdlp;

3
src/bot/oai/mod.rs Normal file
View File

@@ -0,0 +1,3 @@
mod model;
pub use model::*;

160
src/bot/oai/model.rs Normal file
View File

@@ -0,0 +1,160 @@
use serde::{Serialize, Deserialize};
use serde_json::Value;
use siren::ServiceError;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GPTRole {
#[serde(rename = "system")]
System,
#[serde(rename = "user")]
User,
#[serde(rename = "assistant")]
Assistant,
#[serde(rename = "function")]
Function,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionRequest {
pub model: String,
pub messages: Vec<ChatCompletionMessage>,
/// Value between 0 and 2
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f64>,
/// Value between 0 and 1
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub n: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<i64>,
/// Value between -2.0 and 2.0
#[serde(skip_serializing_if = "Option::is_none")]
pub presence_penalty: Option<f64>,
/// Value between -2.0 and 2.0
#[serde(skip_serializing_if = "Option::is_none")]
pub frequency_penalty: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub user: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionMessage {
pub role: GPTRole,
pub content: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionResponse {
pub id: String,
pub object: String,
pub system_fingerprint: Option<String>,
pub created: i64,
pub model: String,
pub usage: Usage,
pub choices: Vec<Choice>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Usage {
pub prompt_tokens: i64,
pub completion_tokens: i64,
pub total_tokens: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Choice {
pub message: ChatCompletionMessage,
pub finish_reason: String,
pub index: i64,
pub logprobs: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
enum ResponseEvent {
ChatCompletionResponse(ChatCompletionResponse),
ResponseError(ResponseError),
// ChatCompletionResponse {
// id: String,
// object: String,
// system_fingerprint: Option<String>,
// created: i64,
// model: String,
// usage: Usage,
// choices: Vec<Choice>,
// },
// ResponseError {
// error: Option<ErrorDetails>,
// message: Option<String>,
// param: Option<String>,
// #[serde(rename = "type")]
// error_type: Option<String>,
// },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ResponseError {
error: Option<ErrorDetails>,
message: Option<String>,
param: Option<String>,
#[serde(rename = "type")]
error_type: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ErrorDetails {
code: Option<String>,
message: Option<String>,
}
pub struct OAI {
pub client: reqwest::Client,
pub base_url: String,
pub service_url: String,
pub max_attempts: i64,
pub token: String,
pub max_tokens: i64,
pub default_model: String,
pub max_context_questions: i64,
}
impl OAI {
pub async fn chat_completion(
&self,
request: ChatCompletionRequest,
) -> Result<ChatCompletionResponse, ServiceError> {
let url = format!("{}/chat/completions", self.base_url);
let response = self
.client
.post(&url)
.bearer_auth(&self.token)
.header("Content-Type", "application/json".to_string())
.json(&request)
.send()
.await;
match response {
Ok(response) => {
let value = response.json::<Value>().await?;
let event: ResponseEvent = serde_json::from_value::<ResponseEvent>(value)?;
match event {
ResponseEvent::ChatCompletionResponse(response) => {
return Ok(response);
},
ResponseEvent::ResponseError(error) => {
return Err(ServiceError {
status: 500,
message: format!("Error: {}", error.message.unwrap()),
});
},
}
}
Err(err) => {
return Err(ServiceError {
status: 500,
message: format!("Error: {}", err),
})
}
}
}
}

40
src/bot/ytdlp/mod.rs Normal file
View File

@@ -0,0 +1,40 @@
mod model;
use std::process::{Child, Command, Output, Stdio};
pub use model::*;
const YOUTUBE_DL_COMMAND: &str = "yt-dlp";
pub struct YtDlp {
command: Command,
args: Vec<String>,
}
impl YtDlp {
pub fn new() -> Self {
let mut cmd = Command::new(YOUTUBE_DL_COMMAND);
cmd
.env("LC_ALL", "en_US.UTF-8")
.stdout(Stdio::piped())
.stdin(Stdio::piped())
.stderr(Stdio::piped());
Self {
command: cmd,
args: Vec::new(),
}
}
pub fn arg(&mut self, arg: &str) -> &mut Self {
self.args.push(arg.to_owned());
self
}
pub fn execute(&mut self) -> std::io::Result<Output> {
self
.command
.args(self.args.clone())
.spawn()
.and_then(Child::wait_with_output)
}
}

10
src/bot/ytdlp/model.rs Normal file
View File

@@ -0,0 +1,10 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Serialize, Deserialize)]
pub struct PlaylistItem {
pub id: String,
pub url: String,
pub title: String,
pub duration: i32,
pub playlist_index: i32,
}