The world's most clever kitty cat
Add /markov
+106
-23
+10
-5
src/brain.rs
+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
+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
+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
+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
+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
+1
-1
src/cmd/weights.rs
+15
-4
src/main.rs
+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
+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")?;