The world's most clever kitty cat

Add /markov

bwc9876.dev a25194c1 a943d036

verified
+106 -23
+10 -5
src/brain.rs
··· 152 152 &self, 153 153 msg: &str, 154 154 is_self: bool, 155 + force_reply: bool, 155 156 mut typing_oneshot: Option<TypingSender>, 156 157 ) -> Option<String> { 157 158 const MAX_TOKENS: usize = 20; ··· 159 160 let mut rng = fastrand::Rng::new(); 160 161 161 162 // Roll if we should reply 162 - if !Self::should_reply(&mut rng, is_self) { 163 + if !force_reply && !Self::should_reply(&mut rng, is_self) { 163 164 debug!("Failed roll"); 164 165 return None; 165 166 } ··· 202 203 typ.send(false).ok(); 203 204 } 204 205 205 - Some(chain.join(" ")) 206 + if chain.is_empty() { 207 + None 208 + } else { 209 + Some(chain.join(" ")).filter(|s| !s.trim().is_empty()) 210 + } 206 211 } 207 212 208 213 pub fn word_count(&self) -> usize { ··· 298 303 hello_edges.0, 299 304 HashMap::from_iter([(Some("world".to_string()), 1)]) 300 305 ); 301 - let reply = brain.respond("hello", false, None); 306 + let reply = brain.respond("hello", false, false, None); 302 307 assert_eq!(reply, Some("world".to_string())); 303 308 } 304 309 ··· 312 317 313 318 for _ in 0..100 { 314 319 // I'm too lazy to mock lazyrand LOL!! 315 - let reply = brain.respond("hello", false, None); 320 + let reply = brain.respond("hello", false, false, None); 316 321 assert_eq!(reply, Some("world".to_string())); 317 322 } 318 323 } ··· 327 332 .collect::<String>(); 328 333 let mut brain = Brain::default(); 329 334 brain.ingest(&msg); 330 - let reply = brain.respond("a", false, None); 335 + let reply = brain.respond("a", false, false, None); 331 336 let expected = LETTERS 332 337 .chars() 333 338 .skip(1)
+5 -2
src/cmd/dump_chain.rs
··· 8 8 9 9 use crate::{ 10 10 BROTLI_BUF_SIZE, BotContext, cmd::DEFER_INTER_RESP_EPHEMERAL, get_brotli_params, prelude::*, 11 + require_owner, 11 12 }; 12 13 13 14 #[derive(CommandModel, CreateCommand)] ··· 19 20 20 21 impl DumpChainCommand { 21 22 pub async fn handle(inter: Interaction, data: CommandData, ctx: Arc<BotContext>) -> Result { 23 + let client = ctx.http.interaction(ctx.app_id); 24 + 25 + require_owner!(inter, ctx, client); 26 + 22 27 let Self { compat } = 23 28 Self::from_interaction(data.into()).context("Failed to parse command data")?; 24 - 25 - let client = ctx.http.interaction(ctx.app_id); 26 29 27 30 client 28 31 .create_response(inter.id, &inter.token, &DEFER_INTER_RESP_EPHEMERAL)
+5 -3
src/cmd/load_chain.rs
··· 8 8 9 9 use crate::{ 10 10 BROTLI_BUF_SIZE, BotContext, brain::Brain, cmd::DEFER_INTER_RESP_EPHEMERAL, prelude::*, 11 - status::update_status, 11 + require_owner, status::update_status, 12 12 }; 13 13 14 14 #[derive(CommandModel, CreateCommand)] ··· 22 22 23 23 impl LoadChainCommand { 24 24 pub async fn handle(inter: Interaction, data: CommandData, ctx: Arc<BotContext>) -> Result { 25 + let client = ctx.http.interaction(ctx.app_id); 26 + 27 + require_owner!(inter, ctx, client); 28 + 25 29 let Self { file, compat } = 26 30 Self::from_interaction(data.into()).context("Failed to parse command data")?; 27 - 28 - let client = ctx.http.interaction(ctx.app_id); 29 31 30 32 client 31 33 .create_response(inter.id, &inter.token, &DEFER_INTER_RESP_EPHEMERAL)
+44
src/cmd/markov.rs
··· 1 + use std::sync::Arc; 2 + 3 + use twilight_interactions::command::{CommandModel, CreateCommand}; 4 + use twilight_model::application::interaction::{Interaction, application_command::CommandData}; 5 + 6 + use crate::{BotContext, cmd::DEFER_INTER_RESP, prelude::*}; 7 + 8 + #[derive(CommandModel, CreateCommand)] 9 + #[command( 10 + name = "markov", 11 + desc = "Trigger a response from bingus! Uses the last word you sent to start the chain" 12 + )] 13 + pub struct MarkovCommand { 14 + /// Prompt bingus should reply to 15 + prompt: String, 16 + } 17 + 18 + impl MarkovCommand { 19 + pub async fn handle(inter: Interaction, data: CommandData, ctx: Arc<BotContext>) -> Result { 20 + let Self { prompt } = 21 + Self::from_interaction(data.into()).context("Failed to parse command data")?; 22 + 23 + let client = ctx.http.interaction(ctx.app_id); 24 + 25 + client 26 + .create_response(inter.id, &inter.token, &DEFER_INTER_RESP) 27 + .await 28 + .context("Failed to defer")?; 29 + 30 + let brain = ctx.brain_handle.read().await; 31 + let content = brain 32 + .respond(&prompt, false, true, None) 33 + .unwrap_or_else(|| String::from("> Bingus couldn't think of what to say!")); 34 + drop(brain); 35 + 36 + client 37 + .update_response(&inter.token) 38 + .content(Some(content.as_str())) 39 + .await 40 + .context("Failed to reply")?; 41 + 42 + Ok(()) 43 + } 44 + }
+24 -2
src/cmd/mod.rs
··· 1 1 mod dump_chain; 2 2 mod load_chain; 3 + mod markov; 3 4 mod weights; 4 5 5 6 use std::sync::Arc; ··· 12 13 InteractionResponse, InteractionResponseData, InteractionResponseType, 13 14 }; 14 15 15 - use crate::cmd::dump_chain::DumpChainCommand; 16 - use crate::cmd::load_chain::LoadChainCommand; 17 16 use crate::{BotContext, prelude::*}; 18 17 18 + use dump_chain::DumpChainCommand; 19 + use load_chain::LoadChainCommand; 20 + use markov::MarkovCommand; 19 21 use weights::WeightsCommand; 20 22 21 23 const DEFER_INTER_RESP: InteractionResponse = InteractionResponse { ··· 40 42 }), 41 43 }; 42 44 45 + #[macro_export] 46 + macro_rules! require_owner { 47 + ($inter:expr, $ctx:expr, $client:expr) => { 48 + if $inter.author_id().is_none_or(|id| !$ctx.owners.contains(&id)) { 49 + let data = twilight_util::builder::InteractionResponseDataBuilder::new() 50 + .content("You're not allowed to run this command!") 51 + .flags(twilight_model::channel::message::MessageFlags::EPHEMERAL) 52 + .build(); 53 + let resp = twilight_model::http::interaction::InteractionResponse { 54 + kind: twilight_model::http::interaction::InteractionResponseType::ChannelMessageWithSource, 55 + data: Some(data), 56 + }; 57 + $client.create_response($inter.id, &$inter.token, &resp).await.context("Failed to deny perms")?; 58 + return Ok(()); 59 + } 60 + }; 61 + } 62 + 43 63 pub async fn register_all_commands(ctx: Arc<BotContext>) -> Result { 44 64 let commands = [ 45 65 WeightsCommand::create_command().into(), 46 66 DumpChainCommand::create_command().into(), 47 67 LoadChainCommand::create_command().into(), 68 + MarkovCommand::create_command().into(), 48 69 ]; 49 70 50 71 let client = ctx.http.interaction(ctx.app_id); ··· 66 87 "weights" => WeightsCommand::handle(inter, data, ctx).await, 67 88 "dump_chain" => DumpChainCommand::handle(inter, data, ctx).await, 68 89 "load_chain" => LoadChainCommand::handle(inter, data, ctx).await, 90 + "markov" => MarkovCommand::handle(inter, data, ctx).await, 69 91 other => { 70 92 warn!("Unknown command send: {other}"); 71 93 Ok(())
+1 -1
src/cmd/weights.rs
··· 11 11 #[derive(CommandModel, CreateCommand)] 12 12 #[command(name = "weights", desc = "Get the weights of a token")] 13 13 pub struct WeightsCommand { 14 - /// Message to send 14 + /// Token to view the weights of 15 15 token: String, 16 16 } 17 17
+15 -4
src/main.rs
··· 37 37 application::interaction::InteractionData, 38 38 id::{ 39 39 Id, 40 - marker::{ApplicationMarker, UserMarker}, 40 + marker::{ApplicationMarker, ChannelMarker, UserMarker}, 41 41 }, 42 42 }; 43 43 ··· 55 55 http: HttpClient, 56 56 self_id: Id<UserMarker>, 57 57 app_id: Id<ApplicationMarker>, 58 + owners: HashSet<Id<UserMarker>>, 58 59 brain_file_path: PathBuf, 59 - reply_channels: HashSet<u64>, 60 + reply_channels: HashSet<Id<ChannelMarker>>, 60 61 brain_handle: BrainHandle, 61 62 shard_sender: MessageSender, 62 63 pending_save: AtomicBool, ··· 137 138 138 139 // Config 139 140 let token_file = std::env::var("TOKEN_FILE").context("Missing TOKEN_FILE env var")?; 140 - let reply_channels: HashSet<u64> = std::env::var("REPLY_CHANNELS") 141 + let reply_channels = std::env::var("REPLY_CHANNELS") 141 142 .context("Missing REPLY_CHANNELS env var")? 142 143 .split(",") 143 - .map(|s| s.trim().parse::<u64>()) 144 + .map(|s| s.trim().parse::<u64>().map(|c| Id::new(c))) 144 145 .collect::<Result<_, _>>() 145 146 .context("Invalid channel IDs for REPLY_CHANNELS")?; 146 147 let brain_file_path = ··· 177 178 178 179 let self_id = app.bot.context("App is not a bot!")?.id; 179 180 181 + let owners = if let Some(user) = app.owner { 182 + HashSet::from_iter([user.id]) 183 + } else if let Some(team) = app.team { 184 + team.members.iter().map(|m| m.user.id).collect() 185 + } else { 186 + warn!("No Owner?? Bingus is free!!!"); 187 + HashSet::new() 188 + }; 189 + 180 190 let context = Arc::new(BotContext { 181 191 http, 182 192 self_id, 183 193 app_id, 194 + owners, 184 195 reply_channels, 185 196 brain_file_path, 186 197 brain_handle,
+2 -6
src/on_message.rs
··· 49 49 }); 50 50 51 51 let brain = ctx.brain_handle.read().await; 52 - if let Some(reply_text) = brain 53 - .respond(msg, is_self, Some(typ_tx)) 54 - .filter(|s| !s.trim().is_empty()) 55 - { 52 + if let Some(reply_text) = brain.respond(msg, is_self, false, Some(typ_tx)) { 56 53 drop(brain); 57 54 done_rx.await.ok(); 58 55 let allowed_mentions = AllowedMentions::default(); ··· 75 72 } 76 73 77 74 pub async fn handle_discord_message(msg: Box<MessageCreate>, ctx: Arc<BotContext>) -> Result { 78 - let channel_id = msg.channel_id.get(); 79 75 let is_self = msg.author.id == ctx.self_id; 80 76 let is_normal_message = matches!(msg.kind, MessageType::Regular | MessageType::Reply); 81 77 let is_ephemeral = msg ··· 89 85 } 90 86 91 87 // Should Reply to Message? 92 - if ctx.reply_channels.contains(&channel_id) { 88 + if ctx.reply_channels.contains(&msg.channel_id) { 93 89 reply_message(&msg.content, msg.id, msg.channel_id, is_self, &ctx) 94 90 .await 95 91 .context("Bingus failed to reply to a message")?;