+229

Cargo.lock

··· 3 3 version = 4 4 4 5 5 [[package]] 6 6 + name = "aho-corasick" 7 7 + version = "1.1.4" 8 8 + source = "registry+https://github.com/rust-lang/crates.io-index" 9 9 + checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" 10 10 + dependencies = [ 11 11 + "memchr", 12 12 + ] 13 13 + 14 14 + [[package]] 6 15 name = "alloc-no-stdlib" 7 16 version = "2.0.4" 8 17 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 18 27 ] 19 28 20 29 [[package]] 30 30 + name = "anstream" 31 31 + version = "0.6.21" 32 32 + source = "registry+https://github.com/rust-lang/crates.io-index" 33 33 + checksum = "43d5b281e737544384e969a5ccad3f1cdd24b48086a0fc1b2a5262a26b8f4f4a" 34 34 + dependencies = [ 35 35 + "anstyle", 36 36 + "anstyle-parse", 37 37 + "anstyle-query", 38 38 + "anstyle-wincon", 39 39 + "colorchoice", 40 40 + "is_terminal_polyfill", 41 41 + "utf8parse", 42 42 + ] 43 43 + 44 44 + [[package]] 45 45 + name = "anstyle" 46 46 + version = "1.0.13" 47 47 + source = "registry+https://github.com/rust-lang/crates.io-index" 48 48 + checksum = "5192cca8006f1fd4f7237516f40fa183bb07f8fbdfedaa0036de5ea9b0b45e78" 49 49 + 50 50 + [[package]] 51 51 + name = "anstyle-parse" 52 52 + version = "0.2.7" 53 53 + source = "registry+https://github.com/rust-lang/crates.io-index" 54 54 + checksum = "4e7644824f0aa2c7b9384579234ef10eb7efb6a0deb83f9630a49594dd9c15c2" 55 55 + dependencies = [ 56 56 + "utf8parse", 57 57 + ] 58 58 + 59 59 + [[package]] 60 60 + name = "anstyle-query" 61 61 + version = "1.1.5" 62 62 + source = "registry+https://github.com/rust-lang/crates.io-index" 63 63 + checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" 64 64 + dependencies = [ 65 65 + "windows-sys 0.61.2", 66 66 + ] 67 67 + 68 68 + [[package]] 69 69 + name = "anstyle-wincon" 70 70 + version = "3.0.11" 71 71 + source = "registry+https://github.com/rust-lang/crates.io-index" 72 72 + checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" 73 73 + dependencies = [ 74 74 + "anstyle", 75 75 + "once_cell_polyfill", 76 76 + "windows-sys 0.61.2", 77 77 + ] 78 78 + 79 79 + [[package]] 21 80 name = "anyhow" 22 81 version = "1.0.102" 23 82 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 68 127 version = "0.1.0" 69 128 dependencies = [ 70 129 "anyhow", 130 130 + "brotli", 131 131 + "colog", 71 132 "fastrand", 133 133 + "log", 72 134 "rmp-serde", 73 135 "rustls", 74 136 "serde", ··· 86 148 version = "2.11.0" 87 149 source = "registry+https://github.com/rust-lang/crates.io-index" 88 150 checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" 151 151 + 152 152 + [[package]] 153 153 + name = "brotli" 154 154 + version = "8.0.2" 155 155 + source = "registry+https://github.com/rust-lang/crates.io-index" 156 156 + checksum = "4bd8b9603c7aa97359dbd97ecf258968c95f3adddd6db2f7e7a5bef101c84560" 157 157 + dependencies = [ 158 158 + "alloc-no-stdlib", 159 159 + "alloc-stdlib", 160 160 + "brotli-decompressor", 161 161 + ] 89 162 90 163 [[package]] 91 164 name = "brotli-decompressor" ··· 137 210 ] 138 211 139 212 [[package]] 213 213 + name = "colog" 214 214 + version = "1.4.0" 215 215 + source = "registry+https://github.com/rust-lang/crates.io-index" 216 216 + checksum = "df62599ba6adc9c6c04a54278c8209125343dc4775f57b9d76c9a4287e58f2bd" 217 217 + dependencies = [ 218 218 + "colored", 219 219 + "env_logger", 220 220 + "log", 221 221 + ] 222 222 + 223 223 + [[package]] 224 224 + name = "colorchoice" 225 225 + version = "1.0.4" 226 226 + source = "registry+https://github.com/rust-lang/crates.io-index" 227 227 + checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" 228 228 + 229 229 + [[package]] 230 230 + name = "colored" 231 231 + version = "3.1.1" 232 232 + source = "registry+https://github.com/rust-lang/crates.io-index" 233 233 + checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" 234 234 + dependencies = [ 235 235 + "windows-sys 0.61.2", 236 236 + ] 237 237 + 238 238 + [[package]] 140 239 name = "combine" 141 240 version = "4.6.7" 142 241 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 198 297 checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" 199 298 200 299 [[package]] 300 300 + name = "env_filter" 301 301 + version = "1.0.0" 302 302 + source = "registry+https://github.com/rust-lang/crates.io-index" 303 303 + checksum = "7a1c3cc8e57274ec99de65301228b537f1e4eedc1b8e0f9411c6caac8ae7308f" 304 304 + dependencies = [ 305 305 + "log", 306 306 + "regex", 307 307 + ] 308 308 + 309 309 + [[package]] 310 310 + name = "env_logger" 311 311 + version = "0.11.9" 312 312 + source = "registry+https://github.com/rust-lang/crates.io-index" 313 313 + checksum = "b2daee4ea451f429a58296525ddf28b45a3b64f1acf6587e2067437bb11e218d" 314 314 + dependencies = [ 315 315 + "anstream", 316 316 + "anstyle", 317 317 + "env_filter", 318 318 + "jiff", 319 319 + "log", 320 320 + ] 321 321 + 322 322 + [[package]] 201 323 name = "equivalent" 202 324 version = "1.0.2" 203 325 source = "registry+https://github.com/rust-lang/crates.io-index" 204 326 checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" 327 327 + 328 328 + [[package]] 329 329 + name = "errno" 330 330 + version = "0.3.14" 331 331 + source = "registry+https://github.com/rust-lang/crates.io-index" 332 332 + checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" 333 333 + dependencies = [ 334 334 + "libc", 335 335 + "windows-sys 0.52.0", 336 336 + ] 205 337 206 338 [[package]] 207 339 name = "fastrand" ··· 428 560 ] 429 561 430 562 [[package]] 563 563 + name = "is_terminal_polyfill" 564 564 + version = "1.70.2" 565 565 + source = "registry+https://github.com/rust-lang/crates.io-index" 566 566 + checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" 567 567 + 568 568 + [[package]] 431 569 name = "itoa" 432 570 version = "1.0.17" 433 571 source = "registry+https://github.com/rust-lang/crates.io-index" 434 572 checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" 435 573 436 574 [[package]] 575 575 + name = "jiff" 576 576 + version = "0.2.23" 577 577 + source = "registry+https://github.com/rust-lang/crates.io-index" 578 578 + checksum = "1a3546dc96b6d42c5f24902af9e2538e82e39ad350b0c766eb3fbf2d8f3d8359" 579 579 + dependencies = [ 580 580 + "jiff-static", 581 581 + "log", 582 582 + "portable-atomic", 583 583 + "portable-atomic-util", 584 584 + "serde_core", 585 585 + ] 586 586 + 587 587 + [[package]] 588 588 + name = "jiff-static" 589 589 + version = "0.2.23" 590 590 + source = "registry+https://github.com/rust-lang/crates.io-index" 591 591 + checksum = "2a8c8b344124222efd714b73bb41f8b5120b27a7cc1c75593a6ff768d9d05aa4" 592 592 + dependencies = [ 593 593 + "proc-macro2", 594 594 + "quote", 595 595 + "syn", 596 596 + ] 597 597 + 598 598 + [[package]] 437 599 name = "jni" 438 600 version = "0.21.1" 439 601 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 525 687 checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" 526 688 527 689 [[package]] 690 690 + name = "once_cell_polyfill" 691 691 + version = "1.70.2" 692 692 + source = "registry+https://github.com/rust-lang/crates.io-index" 693 693 + checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" 694 694 + 695 695 + [[package]] 528 696 name = "openssl-probe" 529 697 version = "0.2.1" 530 698 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 577 745 checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" 578 746 579 747 [[package]] 748 748 + name = "portable-atomic" 749 749 + version = "1.13.1" 750 750 + source = "registry+https://github.com/rust-lang/crates.io-index" 751 751 + checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" 752 752 + 753 753 + [[package]] 754 754 + name = "portable-atomic-util" 755 755 + version = "0.2.5" 756 756 + source = "registry+https://github.com/rust-lang/crates.io-index" 757 757 + checksum = "7a9db96d7fa8782dd8c15ce32ffe8680bbd1e978a43bf51a34d39483540495f5" 758 758 + dependencies = [ 759 759 + "portable-atomic", 760 760 + ] 761 761 + 762 762 + [[package]] 580 763 name = "powerfmt" 581 764 version = "0.2.0" 582 765 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 616 799 ] 617 800 618 801 [[package]] 802 802 + name = "regex" 803 803 + version = "1.12.3" 804 804 + source = "registry+https://github.com/rust-lang/crates.io-index" 805 805 + checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" 806 806 + dependencies = [ 807 807 + "aho-corasick", 808 808 + "memchr", 809 809 + "regex-automata", 810 810 + "regex-syntax", 811 811 + ] 812 812 + 813 813 + [[package]] 814 814 + name = "regex-automata" 815 815 + version = "0.4.14" 816 816 + source = "registry+https://github.com/rust-lang/crates.io-index" 817 817 + checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" 818 818 + dependencies = [ 819 819 + "aho-corasick", 820 820 + "memchr", 821 821 + "regex-syntax", 822 822 + ] 823 823 + 824 824 + [[package]] 825 825 + name = "regex-syntax" 826 826 + version = "0.8.10" 827 827 + source = "registry+https://github.com/rust-lang/crates.io-index" 828 828 + checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" 829 829 + 830 830 + [[package]] 619 831 name = "ring" 620 832 version = "0.17.14" 621 833 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 847 1059 checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" 848 1060 849 1061 [[package]] 1062 1062 + name = "signal-hook-registry" 1063 1063 + version = "1.4.8" 1064 1064 + source = "registry+https://github.com/rust-lang/crates.io-index" 1065 1065 + checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" 1066 1066 + dependencies = [ 1067 1067 + "errno", 1068 1068 + "libc", 1069 1069 + ] 1070 1070 + 1071 1071 + [[package]] 850 1072 name = "simdutf8" 851 1073 version = "0.1.5" 852 1074 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 951 1173 "libc", 952 1174 "mio", 953 1175 "pin-project-lite", 1176 1176 + "signal-hook-registry", 954 1177 "socket2", 955 1178 "tokio-macros", 956 1179 "windows-sys 0.61.2", ··· 1199 1422 version = "0.9.0" 1200 1423 source = "registry+https://github.com/rust-lang/crates.io-index" 1201 1424 checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" 1425 1425 + 1426 1426 + [[package]] 1427 1427 + name = "utf8parse" 1428 1428 + version = "0.2.2" 1429 1429 + source = "registry+https://github.com/rust-lang/crates.io-index" 1430 1430 + checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" 1202 1431 1203 1432 [[package]] 1204 1433 name = "walkdir"

+9 -1

Cargo.toml

··· 5 5 6 6 [dependencies] 7 7 anyhow = "1.0.102" 8 8 + brotli = "8.0.2" 9 9 + colog = "1.4.0" 8 10 fastrand = "2.3.0" 11 11 + log = "0.4.29" 9 12 rmp-serde = "1.3.1" 10 13 rustls = "0.23.37" 11 14 serde = { version = "1.0.228", features = ["derive"] } 12 12 - tokio = { version = "1.50.0", features = ["macros", "rt-multi-thread"] } 15 15 + tokio = { version = "1.50.0", features = [ 16 16 + "macros", 17 17 + "rt-multi-thread", 18 18 + "fs", 19 19 + "signal", 20 20 + ] } 13 21 twilight-cache-inmemory = "0.17.1" 14 22 twilight-gateway = "0.17.1" 15 23 twilight-http = "0.17.1"

+37 -7

src/brain.rs

··· 1 1 + #![allow(unused)] 2 2 + 1 3 use std::collections::HashMap; 2 4 5 5 + use log::debug; 3 6 use serde::{Deserialize, Serialize}; 7 7 + use tokio::sync::oneshot; 4 8 5 9 /// Some = Word, None = End Message 6 10 pub type Token = Option<String>; ··· 11 15 12 16 #[derive(Default, Debug, Clone, Serialize, Deserialize)] 13 17 pub struct Brain(HashMap<Token, Edges>); 18 18 + 19 19 + pub type TypingSender = oneshot::Sender<bool>; 20 20 + pub type TypingReceiver = oneshot::Receiver<bool>; 14 21 15 22 pub fn format_token(tok: &Token) -> String { 16 23 if let Some(w) = tok { ··· 60 67 } 61 68 } 62 69 70 70 + const FORCE_REPLIES: bool = cfg!(test) || (option_env!("BINGUS_FORCE_REPLY").is_some()); 71 71 + 63 72 impl Brain { 64 73 fn normalize_token(word: &str) -> Token { 65 74 let w = if word.starts_with("http://") || word.starts_with("https://") { ··· 84 93 } 85 94 86 95 fn should_reply(rand: &mut fastrand::Rng, is_self: bool) -> bool { 87 87 - let chance = if is_self { 80 } else { 45 }; 96 96 + let chance = if is_self { 45 } else { 80 }; 88 97 let roll = rand.u8(0..=100); 89 98 90 90 - cfg!(test) || roll <= chance 99 99 + (FORCE_REPLIES && !is_self) || roll <= chance 91 100 } 92 101 93 102 fn extract_final_token(msg: &str) -> Option<Token> { ··· 106 115 } 107 116 } 108 117 109 109 - pub fn ingest(&mut self, msg: &str) { 118 118 + pub fn ingest(&mut self, msg: &str) -> bool { 119 119 + let mut learned_new_word = false; 110 120 // This is a silly way to do windows rust ppl :sob: 111 121 let _ = Self::parse(msg) 112 122 .map_windows(|[from, to]| { 113 113 - eprintln!("{from:?} {to:?}"); 114 123 if let Some(edge) = self.0.get_mut(from) { 115 124 edge.increment_token(to); 116 125 } else { 117 126 let new = Edges(HashMap::from_iter([(to.clone(), 1)]), 1); 118 127 self.0.insert(from.clone(), new); 128 128 + learned_new_word = true; 119 129 } 120 130 }) 121 131 .collect::<Vec<_>>(); 132 132 + 133 133 + learned_new_word 122 134 } 123 135 124 136 pub fn merge_from(&mut self, other: Self) { ··· 131 143 } 132 144 } 133 145 134 134 - pub fn respond(&self, msg: &str, is_self: bool) -> Option<String> { 146 146 + pub fn respond( 147 147 + &self, 148 148 + msg: &str, 149 149 + is_self: bool, 150 150 + mut typing_oneshot: Option<TypingSender>, 151 151 + ) -> Option<String> { 135 152 const MAX_TOKENS: usize = 20; 136 153 137 154 let mut rng = fastrand::Rng::new(); 138 155 139 156 // Roll if we should reply 140 157 if !Self::should_reply(&mut rng, is_self) { 158 158 + debug!("Failed roll"); 141 159 return None; 142 160 } 143 161 ··· 147 165 Self::extract_final_token(msg).or_else(|| self.random_token(&mut rng))?; 148 166 149 167 let mut chain = Vec::with_capacity(MAX_TOKENS); 168 168 + let mut has_triggered_typing = false; 150 169 151 170 while let Some(tok) = current_token 152 171 && chain.len() <= MAX_TOKENS ··· 155 174 let next = edges.sample(&mut rng).flatten(); 156 175 if let Some(ref s) = next { 157 176 chain.push(s.clone()); 177 177 + if !has_triggered_typing && let Some(typ) = typing_oneshot.take() { 178 178 + typ.send(true).ok(); 179 179 + } 158 180 } 159 181 current_token = next; 160 182 } else { ··· 162 184 } 163 185 } 164 186 187 187 + if let Some(typ) = typing_oneshot.take() { 188 188 + typ.send(false).ok(); 189 189 + } 190 190 + 165 191 Some(chain.join(" ")) 192 192 + } 193 193 + 194 194 + pub fn word_count(&self) -> usize { 195 195 + self.0.len() 166 196 } 167 197 168 198 pub fn get_weights(&self, tok: &str) -> Option<&Edges> { ··· 211 241 hello_edges.0, 212 242 HashMap::from_iter([(Some("world".to_string()), 1)]) 213 243 ); 214 214 - let reply = brain.respond("hello", false); 244 244 + let reply = brain.respond("hello", false, None); 215 245 assert_eq!(reply, Some("world".to_string())); 216 246 } 217 247 ··· 225 255 .join(" "); 226 256 let mut brain = Brain::default(); 227 257 brain.ingest(&msg); 228 228 - let reply = brain.respond("a", false); 258 258 + let reply = brain.respond("a", false, None); 229 259 let expected = LETTERS 230 260 .chars() 231 261 .skip(1)

+165 -26

src/main.rs

··· 1 1 #![feature(iter_map_windows)] 2 2 - #![allow(unused)] 3 2 4 3 mod brain; 4 4 + mod on_message; 5 5 + mod status; 5 6 6 7 pub mod prelude { 7 8 pub use anyhow::Context; ··· 9 10 pub type Result<T = (), E = anyhow::Error> = StdResult<T, E>; 10 11 } 11 12 12 12 - use std::{collections::HashSet, sync::Arc}; 13 13 + use std::{ 14 14 + collections::HashSet, 15 15 + fs::File, 16 16 + path::{Path, PathBuf}, 17 17 + sync::{ 18 18 + Arc, 19 19 + atomic::{AtomicBool, Ordering}, 20 20 + }, 21 21 + }; 13 22 23 23 + use brotli::enc::BrotliEncoderParams; 24 24 + use log::{debug, error, info, warn}; 14 25 use prelude::*; 26 26 + use tokio::{ 27 27 + sync::Mutex, 28 28 + time::{self, Duration}, 29 29 + }; 15 30 use twilight_cache_inmemory::{DefaultInMemoryCache, ResourceType}; 16 16 - use twilight_gateway::{Event, EventTypeFlags, Intents, Shard, ShardId, StreamExt}; 31 31 + use twilight_gateway::{ 32 32 + CloseFrame, Event, EventTypeFlags, Intents, MessageSender, Shard, ShardId, StreamExt, 33 33 + }; 17 34 use twilight_http::Client as HttpClient; 35 35 + use twilight_model::id::{Id, marker::UserMarker}; 36 36 + 37 37 + use crate::{brain::Brain, on_message::handle_discord_message, status::update_status}; 38 38 + 39 39 + pub type BrainHandle = Mutex<Brain>; 18 40 19 41 #[derive(Debug)] 20 20 - struct BotContext { 42 42 + pub struct BotContext { 21 43 http: HttpClient, 22 22 - reply_channels: HashSet<String>, 44 44 + self_id: Id<UserMarker>, 45 45 + brain_file_path: PathBuf, 46 46 + reply_channels: HashSet<u64>, 47 47 + brain_handle: BrainHandle, 48 48 + shard_sender: MessageSender, 49 49 + pending_save: AtomicBool, 23 50 } 24 51 25 25 - async fn handle_discord_event(event: Event, _ctx: Arc<BotContext>) -> Result { 52 52 + async fn handle_discord_event(event: Event, ctx: Arc<BotContext>) -> Result { 26 53 match event { 27 27 - Event::MessageCreate(msg) => { 28 28 - let channel_id = msg.channel_id.to_string(); 29 29 - eprintln!("id: {channel_id}"); 30 30 - } 54 54 + Event::MessageCreate(msg) => handle_discord_message(msg, ctx).await?, 31 55 Event::Ready(ev) => { 32 32 - eprintln!("Connected to gateway as {}", ev.user.name); 56 56 + info!("Connected to gateway as {}", ev.user.name); 57 57 + let brain = ctx.brain_handle.lock().await; 58 58 + update_status(&*brain, &ctx.shard_sender).context("Failed to update status")?; 33 59 } 34 34 - _ => {} 60 60 + _ => { 61 61 + debug!("Ev: {event:?}"); 62 62 + } 35 63 } 36 64 37 65 Ok(()) 38 66 } 39 67 68 68 + fn load_brain(path: &Path) -> Result<Option<Brain>> { 69 69 + if path.exists() { 70 70 + let mut file = File::open(path).context("Failed to open brain file")?; 71 71 + let mut brotli_stream = brotli::Decompressor::new(&mut file, 4096); 72 72 + rmp_serde::from_read(&mut brotli_stream) 73 73 + .map(|b| Some(b)) 74 74 + .context("Failed to decode brain file") 75 75 + } else { 76 76 + Ok(None) 77 77 + } 78 78 + } 79 79 + 80 80 + async fn save_brain(ctx: Arc<BotContext>) -> Result { 81 81 + let mut file = File::create(&ctx.brain_file_path).context("Failed to open brain file")?; 82 82 + let params = BrotliEncoderParams::default(); 83 83 + let mut brotli_writer = brotli::CompressorWriter::with_params(&mut file, 4096, &params); 84 84 + let brain = ctx.brain_handle.lock().await; 85 85 + rmp_serde::encode::write(&mut brotli_writer, &*brain) 86 86 + .context("Failed to write serialized brain")?; 87 87 + debug!("Saved brain file"); 88 88 + Ok(()) 89 89 + } 90 90 + 40 91 #[tokio::main] 41 92 async fn main() -> Result { 93 93 + let mut clog = colog::default_builder(); 94 94 + clog.filter( 95 95 + None, 96 96 + if cfg!(debug_assertions) { 97 97 + log::LevelFilter::Debug 98 98 + } else { 99 99 + log::LevelFilter::Info 100 100 + }, 101 101 + ); 102 102 + clog.try_init().context("Failed to initialize colog")?; 103 103 + 104 104 + info!("Start of bingus-bot {}", env!("CARGO_PKG_VERSION")); 105 105 + 42 106 // Config 43 107 let token_file = std::env::var("TOKEN_FILE").context("Missing TOKEN_FILE env var")?; 44 44 - let reply_channels: HashSet<String> = HashSet::from_iter( 45 45 - std::env::var("REPLY_CHANNELS") 46 46 - .context("Missing REPLY_CHANNELS env var")? 47 47 - .split(",") 48 48 - .map(|s| s.trim().to_string()), 49 49 - ); 108 108 + let reply_channels: HashSet<u64> = std::env::var("REPLY_CHANNELS") 109 109 + .context("Missing REPLY_CHANNELS env var")? 110 110 + .split(",") 111 111 + .map(|s| s.trim().parse::<u64>()) 112 112 + .collect::<Result<_, _>>() 113 113 + .context("Invalid channel IDs for REPLY_CHANNELS")?; 114 114 + let brain_file_path = 115 115 + PathBuf::from(std::env::var("BRAIN_FILE").unwrap_or_else(|_| "brain.msgpackz".to_string())); 50 116 let intents = Intents::GUILD_MESSAGES | Intents::MESSAGE_CONTENT; 51 117 52 118 // Read token 53 119 let token = std::fs::read_to_string(token_file).context("Failed to read bot token")?; 54 120 let token = token.trim(); 55 121 122 122 + // Read Brain 123 123 + let brain = if let Some(brain) = load_brain(&brain_file_path)? { 124 124 + info!("Loading brain from {brain_file_path:?}"); 125 125 + brain 126 126 + } else { 127 127 + info!("Creating new brain file at {brain_file_path:?}"); 128 128 + Brain::default() 129 129 + }; 130 130 + let brain_handle = Mutex::new(brain); 131 131 + 56 132 // Init 57 133 let mut shard = Shard::new(ShardId::ONE, token.to_string(), intents); 58 134 let http = HttpClient::new(token.to_string()); ··· 65 141 ) 66 142 .build(); 67 143 144 144 + let self_id = http 145 145 + .current_user_application() 146 146 + .await 147 147 + .context("Failed to get current App")? 148 148 + .model() 149 149 + .await 150 150 + .context("Failed to deserialize")? 151 151 + .bot 152 152 + .context("App is not a bot!")? 153 153 + .id; 154 154 + 68 155 let context = Arc::new(BotContext { 69 156 http, 157 157 + self_id, 70 158 reply_channels, 159 159 + brain_file_path, 160 160 + brain_handle, 161 161 + shard_sender: shard.sender(), 162 162 + pending_save: AtomicBool::new(false), 71 163 }); 72 164 73 73 - // Event Loop 74 74 - while let Some(res) = shard.next_event(EventTypeFlags::all()).await { 75 75 - match res { 76 76 - Ok(event) => { 77 77 - cache.update(&event); 78 78 - tokio::spawn(handle_discord_event(event, Arc::clone(&context))); 165 165 + info!("Ensuring brain is writable..."); 166 166 + save_brain(context.clone()) 167 167 + .await 168 168 + .context("Brain file is not writable")?; 169 169 + info!("Brain file saved"); 170 170 + 171 171 + let mut interval = time::interval(Duration::from_secs(60)); 172 172 + interval.tick().await; 173 173 + tokio::pin!(interval); 174 174 + 175 175 + info!("Connecting to gateway..."); 176 176 + 177 177 + loop { 178 178 + tokio::select! { 179 179 + Ok(()) = tokio::signal::ctrl_c() => { 180 180 + info!("SIGINT: Closing connection and saving"); 181 181 + shard.close(CloseFrame::NORMAL); 182 182 + break; 79 183 } 80 80 - Err(why) => { 81 81 - eprintln!("Failed to receive event: {why:?}"); 184 184 + _ = interval.tick() => { 185 185 + debug!("Save Interval"); 186 186 + if context.pending_save.load(Ordering::Relaxed) { 187 187 + let ctx = context.clone(); 188 188 + tokio::spawn(async move { 189 189 + if let Err(why) = save_brain(ctx.clone()).await { 190 190 + error!("Failed to save brain file:\n{why:?}"); 191 191 + } 192 192 + ctx.pending_save.store(true, Ordering::Relaxed); 193 193 + }); 194 194 + } 195 195 + }, 196 196 + opt = shard.next_event(EventTypeFlags::all()) => { 197 197 + match opt { 198 198 + Some(Ok(event)) => { 199 199 + cache.update(&event); 200 200 + let ctx = context.clone(); 201 201 + tokio::spawn(async move { 202 202 + if let Err(why) = handle_discord_event(event, ctx).await { 203 203 + error!("Error while processing Discord event:\n{why:?}"); 204 204 + } 205 205 + }); 206 206 + } 207 207 + Some(Err(why)) => { 208 208 + warn!("Failed to receive event:\n{why:?}"); 209 209 + } 210 210 + None => { 211 211 + info!("Disconnected from Discord: Saving brain and exiting"); 212 212 + break; 213 213 + } 214 214 + } 82 215 } 83 216 } 84 217 } 218 218 + 219 219 + save_brain(context) 220 220 + .await 221 221 + .context("Failed to write brain file on exit")?; 222 222 + 223 223 + info!("Save Complete, Exiting"); 85 224 86 225 Ok(()) 87 226 }

+72

src/on_message.rs

··· 1 1 + use std::{ 2 2 + boxed::Box, 3 3 + sync::{Arc, atomic::Ordering}, 4 4 + }; 5 5 + 6 6 + use log::warn; 7 7 + use twilight_model::{ 8 8 + channel::message::{AllowedMentions, MessageFlags, MessageType}, 9 9 + gateway::payload::incoming::MessageCreate, 10 10 + }; 11 11 + 12 12 + use crate::{BotContext, prelude::*, status::update_status}; 13 13 + 14 14 + pub async fn handle_discord_message(msg: Box<MessageCreate>, ctx: Arc<BotContext>) -> Result { 15 15 + let channel_id = msg.channel_id.get(); 16 16 + let is_self = msg.author.id == ctx.self_id; 17 17 + let is_normal_message = matches!(msg.kind, MessageType::Regular | MessageType::Reply); 18 18 + let is_ephemeral = msg 19 19 + .flags 20 20 + .is_some_and(|flags| flags.contains(MessageFlags::EPHEMERAL)); 21 21 + let is_dm = msg.guild_id.is_none(); 22 22 + 23 23 + // Should Ingest Message? 24 24 + if is_self || !is_normal_message || is_ephemeral || is_dm { 25 25 + return Ok(()); 26 26 + } 27 27 + 28 28 + let mut brain = ctx.brain_handle.lock().await; 29 29 + let learned_new_word = brain.ingest(&msg.content); 30 30 + ctx.pending_save.store(true, Ordering::Relaxed); 31 31 + 32 32 + if learned_new_word { 33 33 + update_status(&*brain, &ctx.shard_sender).context("Failed to update status")?; 34 34 + } 35 35 + 36 36 + // Should Reply to Message? 37 37 + if !ctx.reply_channels.contains(&channel_id) { 38 38 + return Ok(()); 39 39 + } 40 40 + 41 41 + let (typ_tx, typ_rx) = tokio::sync::oneshot::channel(); 42 42 + let (done_tx, done_rx) = tokio::sync::oneshot::channel(); 43 43 + 44 44 + let ctx_typ = ctx.clone(); 45 45 + let typ_id = msg.channel_id; 46 46 + tokio::spawn(async move { 47 47 + if typ_rx.await.ok().is_some_and(|start| start) { 48 48 + if let Err(why) = ctx_typ.http.create_typing_trigger(typ_id).await { 49 49 + warn!("Failed to set typing indicator:\n{why:?}"); 50 50 + } 51 51 + } 52 52 + done_tx.send(()).ok(); 53 53 + }); 54 54 + 55 55 + if let Some(reply_text) = brain 56 56 + .respond(&msg.content, is_self, Some(typ_tx)) 57 57 + .filter(|s| !s.trim().is_empty()) 58 58 + { 59 59 + drop(brain); 60 60 + done_rx.await.ok(); 61 61 + ctx.http 62 62 + .create_message(msg.channel_id) 63 63 + .content(&reply_text) 64 64 + .reply(msg.id) 65 65 + .fail_if_not_exists(false) 66 66 + .allowed_mentions(Some(&AllowedMentions::default())) 67 67 + .await 68 68 + .context("Failed to send message")?; 69 69 + } 70 70 + 71 71 + Ok(()) 72 72 + }

+41

src/status.rs

··· 1 1 + use log::debug; 2 2 + use twilight_gateway::MessageSender; 3 3 + use twilight_model::gateway::{ 4 4 + payload::outgoing::UpdatePresence, 5 5 + presence::{Activity, ActivityType, Status}, 6 6 + }; 7 7 + 8 8 + use crate::{brain::Brain, prelude::*}; 9 9 + 10 10 + pub fn update_status(brain: &Brain, sender: &MessageSender) -> Result { 11 11 + let words = brain.word_count(); 12 12 + 13 13 + let activity = Activity { 14 14 + application_id: None, 15 15 + assets: None, 16 16 + buttons: Vec::new(), 17 17 + created_at: None, 18 18 + details: None, 19 19 + emoji: None, 20 20 + flags: None, 21 21 + id: None, 22 22 + instance: None, 23 23 + kind: ActivityType::Custom, 24 24 + name: "Bingus".to_string(), 25 25 + party: None, 26 26 + secrets: None, 27 27 + state: Some(format!("I know {words} words!")), 28 28 + timestamps: None, 29 29 + url: None, 30 30 + }; 31 31 + 32 32 + let status = UpdatePresence::new(vec![activity], false, None, Status::Online) 33 33 + .context("Failed to make status")?; 34 34 + 35 35 + sender 36 36 + .command(&status) 37 37 + .context("Failed to send to gateway")?; 38 38 + 39 39 + debug!("Sent status update"); 40 40 + Ok(()) 41 41 + }