The world's most clever kitty cat
Add /weights command
+158
-19
+4
-4
src/brain.rs
+4
-4
src/brain.rs
···
251
251
let msg = LETTERS
252
252
.chars()
253
253
.map(|c| c.to_string())
254
-
.collect::<Vec<_>>()
255
-
.join(" ");
254
+
.intersperse(" ".to_string())
255
+
.collect::<String>();
256
256
let mut brain = Brain::default();
257
257
brain.ingest(&msg);
258
258
let reply = brain.respond("a", false, None);
···
261
261
.skip(1)
262
262
.take(21)
263
263
.map(|c| c.to_string())
264
-
.collect::<Vec<_>>()
265
-
.join(" ");
264
+
.intersperse(" ".to_string())
265
+
.collect::<String>();
266
266
assert_eq!(reply, Some(expected));
267
267
}
268
268
+44
src/cmd/mod.rs
+44
src/cmd/mod.rs
···
1
+
mod weights;
2
+
3
+
use std::sync::Arc;
4
+
5
+
use log::warn;
6
+
use twilight_interactions::command::CreateCommand;
7
+
use twilight_model::application::interaction::{Interaction, application_command::CommandData};
8
+
use twilight_model::http::interaction::{InteractionResponse, InteractionResponseType};
9
+
10
+
use crate::{BotContext, prelude::*};
11
+
12
+
use weights::WeightsCommand;
13
+
14
+
const DEFER_INTER_RESP: InteractionResponse = InteractionResponse {
15
+
kind: InteractionResponseType::DeferredChannelMessageWithSource,
16
+
data: None,
17
+
};
18
+
19
+
pub async fn register_all_commands(ctx: Arc<BotContext>) -> Result {
20
+
let commands = [WeightsCommand::create_command().into()];
21
+
22
+
let client = ctx.http.interaction(ctx.app_id);
23
+
24
+
client
25
+
.set_global_commands(&commands)
26
+
.await
27
+
.context("Failed to register app commands")?;
28
+
29
+
Ok(())
30
+
}
31
+
32
+
pub async fn handle_app_command(
33
+
data: CommandData,
34
+
ctx: Arc<BotContext>,
35
+
inter: Interaction,
36
+
) -> Result {
37
+
match &*data.name {
38
+
"weights" => WeightsCommand::handle(inter, data, ctx).await,
39
+
other => {
40
+
warn!("Unknown command send: {other}");
41
+
Ok(())
42
+
}
43
+
}
44
+
}
+70
src/cmd/weights.rs
+70
src/cmd/weights.rs
···
1
+
use std::sync::Arc;
2
+
3
+
use twilight_interactions::command::{CommandModel, CreateCommand};
4
+
use twilight_model::{
5
+
application::interaction::{Interaction, application_command::CommandData},
6
+
http::attachment::Attachment,
7
+
};
8
+
9
+
use crate::{BotContext, BrainHandle, brain::format_token, cmd::DEFER_INTER_RESP, prelude::*};
10
+
11
+
#[derive(CommandModel, CreateCommand)]
12
+
#[command(name = "weights", desc = "Get the weights of a token")]
13
+
pub struct WeightsCommand {
14
+
/// Message to send
15
+
token: String,
16
+
}
17
+
18
+
async fn get_output(token: &str, brain: &BrainHandle) -> Option<String> {
19
+
let brain = brain.lock().await;
20
+
21
+
brain.get_weights(token).map(|edges| {
22
+
let sep = String::from("\n");
23
+
let all_weights = edges
24
+
.iter_weights()
25
+
.map(|(token, weight, chance)| {
26
+
let token_fmt = format_token(token);
27
+
format!("- **{token_fmt}**: {:.1}% ({weight})", chance * 100.0)
28
+
})
29
+
.intersperse(sep)
30
+
.collect::<String>();
31
+
32
+
format!("Weights for {token}:\n{all_weights}")
33
+
})
34
+
}
35
+
36
+
impl WeightsCommand {
37
+
pub async fn handle(inter: Interaction, data: CommandData, ctx: Arc<BotContext>) -> Result {
38
+
let WeightsCommand { token } = WeightsCommand::from_interaction(data.into())
39
+
.context("Failed to parse command data")?;
40
+
41
+
let client = ctx.http.interaction(ctx.app_id);
42
+
43
+
client
44
+
.create_response(inter.id, &inter.token, &DEFER_INTER_RESP)
45
+
.await
46
+
.context("Failed to defer")?;
47
+
48
+
let content = get_output(&token, &ctx.brain_handle)
49
+
.await
50
+
.unwrap_or_else(|| String::from("Bingus doesn't know that word!"));
51
+
52
+
let update = client.update_response(&inter.token);
53
+
54
+
if content.encode_utf16().count() < 2000 {
55
+
update.content(Some(content.as_str())).await
56
+
} else {
57
+
let data = content.into_bytes();
58
+
let attachment = Attachment::from_bytes(String::from("weights.txt"), data, 1);
59
+
update
60
+
.content(Some(
61
+
"Weights were too long to fit into one message, check the text file!",
62
+
))
63
+
.attachments(&[attachment])
64
+
.await
65
+
}
66
+
.context("Failed to reply")?;
67
+
68
+
Ok(())
69
+
}
70
+
}
+40
-15
src/main.rs
+40
-15
src/main.rs
···
1
1
#![feature(iter_map_windows)]
2
+
#![feature(iter_intersperse)]
2
3
3
4
mod brain;
5
+
mod cmd;
4
6
mod on_message;
5
7
mod status;
6
8
···
31
33
CloseFrame, Event, EventTypeFlags, Intents, MessageSender, Shard, ShardId, StreamExt,
32
34
};
33
35
use twilight_http::Client as HttpClient;
34
-
use twilight_model::id::{Id, marker::UserMarker};
36
+
use twilight_model::{
37
+
application::interaction::InteractionData,
38
+
id::{
39
+
Id,
40
+
marker::{ApplicationMarker, UserMarker},
41
+
},
42
+
};
35
43
36
-
use crate::{brain::Brain, on_message::handle_discord_message, status::update_status};
44
+
use crate::{
45
+
brain::Brain,
46
+
cmd::{handle_app_command, register_all_commands},
47
+
on_message::handle_discord_message,
48
+
status::update_status,
49
+
};
37
50
38
51
pub type BrainHandle = Mutex<Brain>;
39
52
···
41
54
pub struct BotContext {
42
55
http: HttpClient,
43
56
self_id: Id<UserMarker>,
57
+
app_id: Id<ApplicationMarker>,
44
58
brain_file_path: PathBuf,
45
59
reply_channels: HashSet<u64>,
46
60
brain_handle: BrainHandle,
···
51
65
async fn handle_discord_event(event: Event, ctx: Arc<BotContext>) -> Result {
52
66
match event {
53
67
Event::MessageCreate(msg) => handle_discord_message(msg, ctx).await,
68
+
Event::InteractionCreate(mut inter) => {
69
+
if let Some(InteractionData::ApplicationCommand(data)) =
70
+
std::mem::take(&mut inter.0.data)
71
+
{
72
+
handle_app_command(*data, ctx, inter.0).await
73
+
} else {
74
+
Ok(())
75
+
}
76
+
}
54
77
Event::Ready(ev) => {
55
78
info!("Connected to gateway as {}", ev.user.name);
56
79
let brain = ctx.brain_handle.lock().await;
57
80
update_status(&*brain, &ctx.shard_sender).context("Failed to update status")
58
81
}
59
-
_ => {
60
-
debug!("Ev: {event:?}");
61
-
Ok(())
62
-
}
82
+
_ => Ok(()),
63
83
}
64
84
}
65
85
···
131
151
let mut shard = Shard::new(ShardId::ONE, token.to_string(), intents);
132
152
let http = HttpClient::new(token.to_string());
133
153
134
-
let self_id = http
154
+
let app = http
135
155
.current_user_application()
136
156
.await
137
157
.context("Failed to get current App")?
138
158
.model()
139
159
.await
140
-
.context("Failed to deserialize")?
141
-
.bot
142
-
.context("App is not a bot!")?
143
-
.id;
160
+
.context("Failed to deserialize")?;
161
+
162
+
let app_id = app.id;
163
+
164
+
let self_id = app.bot.context("App is not a bot!")?.id;
144
165
145
166
let context = Arc::new(BotContext {
146
167
http,
147
168
self_id,
169
+
app_id,
148
170
reply_channels,
149
171
brain_file_path,
150
172
brain_handle,
···
157
179
.await
158
180
.context("Brain file is not writable")?;
159
181
info!("Brain file saved");
182
+
183
+
info!("Registering Commands...");
184
+
register_all_commands(context.clone()).await?;
160
185
161
186
let mut interval = time::interval(Duration::from_secs(60));
162
187
interval.tick().await;
···
188
213
},
189
214
opt = shard.next_event(EventTypeFlags::all()) => {
190
215
match opt {
216
+
Some(Ok(Event::GatewayClose(_))) | None => {
217
+
info!("Disconnected from Discord: Saving brain and exiting");
218
+
break;
219
+
}
191
220
Some(Ok(event)) => {
192
221
let ctx = context.clone();
193
222
tokio::spawn(async move {
···
198
227
}
199
228
Some(Err(why)) => {
200
229
warn!("Failed to receive event:\n{why:?}");
201
-
}
202
-
None => {
203
-
info!("Disconnected from Discord: Saving brain and exiting");
204
-
break;
205
230
}
206
231
}
207
232
}