+126 -2

Cargo.lock

reviewed

··· 36 36 checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" 37 37 38 38 [[package]] 39 39 + name = "aws-lc-rs" 40 40 + version = "1.16.0" 41 41 + source = "registry+https://github.com/rust-lang/crates.io-index" 42 42 + checksum = "d9a7b350e3bb1767102698302bc37256cbd48422809984b98d292c40e2579aa9" 43 43 + dependencies = [ 44 44 + "aws-lc-sys", 45 45 + "zeroize", 46 46 + ] 47 47 + 48 48 + [[package]] 49 49 + name = "aws-lc-sys" 50 50 + version = "0.37.1" 51 51 + source = "registry+https://github.com/rust-lang/crates.io-index" 52 52 + checksum = "b092fe214090261288111db7a2b2c2118e5a7f30dc2569f1732c4069a6840549" 53 53 + dependencies = [ 54 54 + "cc", 55 55 + "cmake", 56 56 + "dunce", 57 57 + "fs_extra", 58 58 + ] 59 59 + 60 60 + [[package]] 39 61 name = "base64" 40 62 version = "0.22.1" 41 63 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 46 68 version = "0.1.0" 47 69 dependencies = [ 48 70 "anyhow", 71 71 + "fastrand", 49 72 "rmp-serde", 73 73 + "rustls", 50 74 "serde", 51 75 "tokio", 76 76 + "twilight-cache-inmemory", 52 77 "twilight-gateway", 53 78 "twilight-http", 54 79 "twilight-interactions", ··· 103 128 checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" 104 129 105 130 [[package]] 131 131 + name = "cmake" 132 132 + version = "0.1.57" 133 133 + source = "registry+https://github.com/rust-lang/crates.io-index" 134 134 + checksum = "75443c44cd6b379beb8c5b45d85d0773baf31cce901fe7bb252f4eff3008ef7d" 135 135 + dependencies = [ 136 136 + "cc", 137 137 + ] 138 138 + 139 139 + [[package]] 106 140 name = "combine" 107 141 version = "4.6.7" 108 142 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 129 163 checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" 130 164 131 165 [[package]] 166 166 + name = "crossbeam-utils" 167 167 + version = "0.8.21" 168 168 + source = "registry+https://github.com/rust-lang/crates.io-index" 169 169 + checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" 170 170 + 171 171 + [[package]] 172 172 + name = "dashmap" 173 173 + version = "6.1.0" 174 174 + source = "registry+https://github.com/rust-lang/crates.io-index" 175 175 + checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" 176 176 + dependencies = [ 177 177 + "cfg-if", 178 178 + "crossbeam-utils", 179 179 + "hashbrown 0.14.5", 180 180 + "lock_api", 181 181 + "once_cell", 182 182 + "parking_lot_core", 183 183 + ] 184 184 + 185 185 + [[package]] 132 186 name = "deranged" 133 187 version = "0.5.8" 134 188 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 136 190 dependencies = [ 137 191 "powerfmt", 138 192 ] 193 193 + 194 194 + [[package]] 195 195 + name = "dunce" 196 196 + version = "1.0.5" 197 197 + source = "registry+https://github.com/rust-lang/crates.io-index" 198 198 + checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" 139 199 140 200 [[package]] 141 201 name = "equivalent" ··· 162 222 checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" 163 223 164 224 [[package]] 225 225 + name = "fs_extra" 226 226 + version = "1.3.0" 227 227 + source = "registry+https://github.com/rust-lang/crates.io-index" 228 228 + checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" 229 229 + 230 230 + [[package]] 165 231 name = "futures-channel" 166 232 version = "0.3.32" 167 233 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 240 306 "tokio-util", 241 307 "tracing", 242 308 ] 309 309 + 310 310 + [[package]] 311 311 + name = "hashbrown" 312 312 + version = "0.14.5" 313 313 + source = "registry+https://github.com/rust-lang/crates.io-index" 314 314 + checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" 243 315 244 316 [[package]] 245 317 name = "hashbrown" ··· 352 424 checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" 353 425 dependencies = [ 354 426 "equivalent", 355 355 - "hashbrown", 427 427 + "hashbrown 0.16.1", 356 428 ] 357 429 358 430 [[package]] ··· 400 472 checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" 401 473 402 474 [[package]] 475 475 + name = "lock_api" 476 476 + version = "0.4.14" 477 477 + source = "registry+https://github.com/rust-lang/crates.io-index" 478 478 + checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" 479 479 + dependencies = [ 480 480 + "scopeguard", 481 481 + ] 482 482 + 483 483 + [[package]] 403 484 name = "log" 404 485 version = "0.4.29" 405 486 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 456 537 checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" 457 538 dependencies = [ 458 539 "num-traits", 540 540 + ] 541 541 + 542 542 + [[package]] 543 543 + name = "parking_lot_core" 544 544 + version = "0.9.12" 545 545 + source = "registry+https://github.com/rust-lang/crates.io-index" 546 546 + checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" 547 547 + dependencies = [ 548 548 + "cfg-if", 549 549 + "libc", 550 550 + "redox_syscall", 551 551 + "smallvec", 552 552 + "windows-link", 459 553 ] 460 554 461 555 [[package]] ··· 513 607 checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" 514 608 515 609 [[package]] 610 610 + name = "redox_syscall" 611 611 + version = "0.5.18" 612 612 + source = "registry+https://github.com/rust-lang/crates.io-index" 613 613 + checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" 614 614 + dependencies = [ 615 615 + "bitflags", 616 616 + ] 617 617 + 618 618 + [[package]] 516 619 name = "ring" 517 620 version = "0.17.14" 518 621 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 551 654 source = "registry+https://github.com/rust-lang/crates.io-index" 552 655 checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4" 553 656 dependencies = [ 657 657 + "aws-lc-rs", 658 658 + "log", 554 659 "once_cell", 555 660 "rustls-pki-types", 556 661 "rustls-webpki", ··· 612 717 source = "registry+https://github.com/rust-lang/crates.io-index" 613 718 checksum = "d7df23109aa6c1567d1c575b9952556388da57401e4ace1d15f79eedad0d8f53" 614 719 dependencies = [ 720 720 + "aws-lc-rs", 615 721 "ring", 616 722 "rustls-pki-types", 617 723 "untrusted", ··· 634 740 dependencies = [ 635 741 "windows-sys 0.61.2", 636 742 ] 743 743 + 744 744 + [[package]] 745 745 + name = "scopeguard" 746 746 + version = "1.2.0" 747 747 + source = "registry+https://github.com/rust-lang/crates.io-index" 748 748 + checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" 637 749 638 750 [[package]] 639 751 name = "security-framework" ··· 945 1057 checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" 946 1058 947 1059 [[package]] 1060 1060 + name = "twilight-cache-inmemory" 1061 1061 + version = "0.17.1" 1062 1062 + source = "registry+https://github.com/rust-lang/crates.io-index" 1063 1063 + checksum = "3d205ec8d1fc62db874cc7af787cc5919a6d8513974159de950764917c563f30" 1064 1064 + dependencies = [ 1065 1065 + "bitflags", 1066 1066 + "dashmap", 1067 1067 + "serde", 1068 1068 + "twilight-model", 1069 1069 + ] 1070 1070 + 1071 1071 + [[package]] 948 1072 name = "twilight-gateway" 949 1073 version = "0.17.1" 950 1074 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 1005 1129 source = "registry+https://github.com/rust-lang/crates.io-index" 1006 1130 checksum = "0515b0c30814068a7540fcb5f58b634259ca453fa335d42c3b2c8f2b06ac6a59" 1007 1131 dependencies = [ 1008 1008 - "hashbrown", 1132 1132 + "hashbrown 0.16.1", 1009 1133 "tokio", 1010 1134 "tokio-util", 1011 1135 "tracing",

+3

Cargo.toml

reviewed

··· 5 5 6 6 [dependencies] 7 7 anyhow = "1.0.102" 8 8 + fastrand = "2.3.0" 8 9 rmp-serde = "1.3.1" 10 10 + rustls = "0.23.37" 9 11 serde = { version = "1.0.228", features = ["derive"] } 10 12 tokio = { version = "1.50.0", features = ["macros", "rt-multi-thread"] } 13 13 + twilight-cache-inmemory = "0.17.1" 11 14 twilight-gateway = "0.17.1" 12 15 twilight-http = "0.17.1" 13 16 twilight-interactions = "0.17.0"

+270

src/brain.rs

reviewed

··· 1 1 + use std::collections::HashMap; 2 2 + 3 3 + use serde::{Deserialize, Serialize}; 4 4 + 5 5 + 6 6 + /// Some = Word, None = End Message 7 7 + pub type Token = Option<String>; 8 8 + pub type Weight = u16; 9 9 + 10 10 + #[derive(Default, Debug, Clone, Serialize, Deserialize)] 11 11 + pub struct Edges(HashMap<Token, Weight>, u64); 12 12 + 13 13 + #[derive(Default, Debug, Clone, Serialize, Deserialize)] 14 14 + pub struct Brain(HashMap<Token, Edges>); 15 15 + 16 16 + pub fn format_token(tok: &Token) -> String { 17 17 + if let Some(w) = tok { 18 18 + w.clone() 19 19 + } else { 20 20 + "~END".to_string() 21 21 + } 22 22 + } 23 23 + 24 24 + impl Edges { 25 25 + fn increment_token(&mut self, tok: &Token) { 26 26 + if let Some(w) = self.0.get_mut(tok) { 27 27 + *w = w.saturating_add(1); 28 28 + } else { 29 29 + self.0.insert(tok.clone(), 1); 30 30 + } 31 31 + self.1 = self.1.saturating_add(1); 32 32 + } 33 33 + 34 34 + fn merge_from(&mut self, other: Self) { 35 35 + self.0.reserve(other.0.len()); 36 36 + for (k, v) in other.0.into_iter() { 37 37 + if let Some(w) = self.0.get_mut(&k) { 38 38 + *w = w.saturating_add(v); 39 39 + } else { 40 40 + self.0.insert(k, v); 41 41 + } 42 42 + self.1 = self.1.saturating_add(v as u64); 43 43 + } 44 44 + } 45 45 + 46 46 + fn sample(&self, rand: &mut fastrand::Rng) -> Option<Token> { 47 47 + let mut dist_left = rand.f64() * self.1 as f64; 48 48 + for (tok, weight) in self.0.iter() { 49 49 + dist_left -= *weight as f64; 50 50 + if dist_left < 0.0 { 51 51 + return Some(tok.clone()); 52 52 + } 53 53 + } 54 54 + None 55 55 + } 56 56 + 57 57 + pub fn iter_weights(&self) -> impl Iterator<Item = (&Token, Weight, f64)> { 58 58 + self.0 59 59 + .iter() 60 60 + .map(|(k, v)| (k, *v, (*v as f64) / (self.1 as f64))) 61 61 + } 62 62 + } 63 63 + 64 64 + impl Brain { 65 65 + fn normalize_token(word: &str) -> Token { 66 66 + let w = if word.starts_with("http://") || word.starts_with("https://") { 67 67 + word.to_string() 68 68 + } else { 69 69 + word.to_ascii_lowercase() 70 70 + }; 71 71 + Some(w) 72 72 + } 73 73 + 74 74 + fn parse(msg: &str) -> impl Iterator<Item = Token> { 75 75 + msg.split_ascii_whitespace() 76 76 + .filter_map(|w| { 77 77 + // Filter out pings, they can get annoying 78 78 + if w.starts_with("<@") && w.ends_with(">") { 79 79 + None 80 80 + } else { 81 81 + Some(Self::normalize_token(w)) 82 82 + } 83 83 + }) 84 84 + .chain(std::iter::once(None)) 85 85 + } 86 86 + 87 87 + fn should_reply(rand: &mut fastrand::Rng, is_self: bool) -> bool { 88 88 + let chance = if is_self { 80 } else { 45 }; 89 89 + let roll = rand.u8(0..=100); 90 90 + 91 91 + cfg!(test) || roll <= chance 92 92 + } 93 93 + 94 94 + fn extract_final_token(msg: &str) -> Option<Token> { 95 95 + msg.split_ascii_whitespace() 96 96 + .last() 97 97 + .map(Self::normalize_token) 98 98 + } 99 99 + 100 100 + fn random_token(&self, rand: &mut fastrand::Rng) -> Option<Token> { 101 101 + let len = self.0.len(); 102 102 + if len == 0 { 103 103 + None 104 104 + } else { 105 105 + let i = rand.usize(..len); 106 106 + self.0.keys().nth(i).cloned() 107 107 + } 108 108 + } 109 109 + 110 110 + pub fn ingest(&mut self, msg: &str) { 111 111 + // This is a silly way to do windows rust ppl :sob: 112 112 + let _ = Self::parse(msg) 113 113 + .map_windows(|[from, to]| { 114 114 + eprintln!("{from:?} {to:?}"); 115 115 + if let Some(edge) = self.0.get_mut(from) { 116 116 + edge.increment_token(to); 117 117 + } else { 118 118 + let new = Edges(HashMap::from_iter([(to.clone(), 1)]), 1); 119 119 + self.0.insert(from.clone(), new); 120 120 + } 121 121 + }) 122 122 + .collect::<Vec<_>>(); 123 123 + } 124 124 + 125 125 + pub fn merge_from(&mut self, other: Self) { 126 126 + for (k, v) in other.0.into_iter() { 127 127 + if let Some(edges) = self.0.get_mut(&k) { 128 128 + edges.merge_from(v); 129 129 + } else { 130 130 + self.0.insert(k, v); 131 131 + } 132 132 + } 133 133 + } 134 134 + 135 135 + pub fn respond(&self, msg: &str, is_self: bool) -> Option<String> { 136 136 + const MAX_TOKENS: usize = 20; 137 137 + 138 138 + let mut rng = fastrand::Rng::new(); 139 139 + 140 140 + // Roll if we should reply 141 141 + if !Self::should_reply(&mut rng, is_self) { 142 142 + return None; 143 143 + } 144 144 + 145 145 + // Get our final token, or a random one if the message has nothing, or don't reply at all 146 146 + // if we have no tokens at all. 147 147 + let mut current_token = 148 148 + Self::extract_final_token(msg).or_else(|| self.random_token(&mut rng))?; 149 149 + 150 150 + let mut chain = Vec::with_capacity(MAX_TOKENS); 151 151 + 152 152 + while let Some(tok) = current_token 153 153 + && chain.len() <= MAX_TOKENS 154 154 + { 155 155 + if let Some(edges) = self.0.get(&Some(tok)) { 156 156 + let next = edges.sample(&mut rng).flatten(); 157 157 + if let Some(ref s) = next { 158 158 + chain.push(s.clone()); 159 159 + } 160 160 + current_token = next; 161 161 + } else { 162 162 + current_token = None; 163 163 + } 164 164 + } 165 165 + 166 166 + Some(chain.join(" ")) 167 167 + } 168 168 + 169 169 + pub fn get_weights(&self, tok: &str) -> Option<&Edges> { 170 170 + self.0.get(&Self::normalize_token(tok)) 171 171 + } 172 172 + } 173 173 + 174 174 + #[cfg(test)] 175 175 + mod tests { 176 176 + use super::*; 177 177 + use std::default::Default; 178 178 + 179 179 + #[test] 180 180 + fn ingest_parse() { 181 181 + let tokens = Brain::parse("Hello world").collect::<Vec<_>>(); 182 182 + assert_eq!( 183 183 + tokens, 184 184 + vec![Some("hello".to_string()), Some("world".to_string()), None] 185 185 + ); 186 186 + } 187 187 + 188 188 + #[test] 189 189 + fn ingest_url() { 190 190 + let tokens = Brain::parse("https://example.com/CAPS-PATH").collect::<Vec<_>>(); 191 191 + assert_eq!( 192 192 + tokens, 193 193 + vec![Some("https://example.com/CAPS-PATH".to_string()), None] 194 194 + ); 195 195 + } 196 196 + 197 197 + #[test] 198 198 + fn ingest_ping() { 199 199 + let tokens = Brain::parse("hi <@1234567>").collect::<Vec<_>>(); 200 200 + assert_eq!(tokens, vec![Some("hi".to_string()), None]); 201 201 + } 202 202 + 203 203 + #[test] 204 204 + fn basic_chain() { 205 205 + let mut brain = Brain::default(); 206 206 + brain.ingest("hello world"); 207 207 + let hello_edges = brain 208 208 + .0 209 209 + .get(&Some("hello".to_string())) 210 210 + .expect("Hello edges not created"); 211 211 + assert_eq!( 212 212 + hello_edges.0, 213 213 + HashMap::from_iter([(Some("world".to_string()), 1)]) 214 214 + ); 215 215 + let reply = brain.respond("hello", false); 216 216 + assert_eq!(reply, Some("world".to_string())); 217 217 + } 218 218 + 219 219 + #[test] 220 220 + fn long_chain() { 221 221 + const LETTERS: &str = "abcdefghijklmnopqrstuvwxyz"; 222 222 + let msg = LETTERS 223 223 + .chars() 224 224 + .map(|c| c.to_string()) 225 225 + .collect::<Vec<_>>() 226 226 + .join(" "); 227 227 + let mut brain = Brain::default(); 228 228 + brain.ingest(&msg); 229 229 + let reply = brain.respond("a", false); 230 230 + let expected = LETTERS 231 231 + .chars() 232 232 + .skip(1) 233 233 + .take(21) 234 234 + .map(|c| c.to_string()) 235 235 + .collect::<Vec<_>>() 236 236 + .join(" "); 237 237 + assert_eq!(reply, Some(expected)); 238 238 + } 239 239 + 240 240 + #[test] 241 241 + fn merge_brain() { 242 242 + let mut brain1 = Brain::default(); 243 243 + let mut brain2 = Brain::default(); 244 244 + 245 245 + brain1.ingest("hello world"); 246 246 + brain2.ingest("hello world"); 247 247 + brain2.ingest("hello world"); 248 248 + brain2.ingest("other word"); 249 249 + 250 250 + brain1.merge_from(brain2); 251 251 + 252 252 + let hello_edges = brain1 253 253 + .0 254 254 + .get(&Some("hello".to_string())) 255 255 + .expect("Hello edges not created"); 256 256 + assert_eq!( 257 257 + hello_edges.0, 258 258 + HashMap::from_iter([(Some("world".to_string()), 3)]) 259 259 + ); 260 260 + 261 261 + let new_edges = brain1 262 262 + .0 263 263 + .get(&Some("other".to_string())) 264 264 + .expect("New edges not created"); 265 265 + assert_eq!( 266 266 + new_edges.0, 267 267 + HashMap::from_iter([(Some("word".to_string()), 1)]) 268 268 + ); 269 269 + } 270 270 + }

+85 -2

src/main.rs

reviewed

··· 1 1 - fn main() { 2 2 - println!("Hello, world!"); 1 1 + #![feature(iter_map_windows)] 2 2 + 3 3 + mod brain; 4 4 + 5 5 + pub mod prelude { 6 6 + pub use anyhow::Context; 7 7 + use std::result::Result as StdResult; 8 8 + pub type Result<T = (), E = anyhow::Error> = StdResult<T, E>; 9 9 + } 10 10 + 11 11 + use std::{collections::HashSet, sync::Arc}; 12 12 + 13 13 + use prelude::*; 14 14 + use twilight_cache_inmemory::{DefaultInMemoryCache, ResourceType}; 15 15 + use twilight_gateway::{Event, EventTypeFlags, Intents, Shard, ShardId, StreamExt}; 16 16 + use twilight_http::Client as HttpClient; 17 17 + 18 18 + #[derive(Debug)] 19 19 + struct BotContext { 20 20 + http: HttpClient, 21 21 + reply_channels: HashSet<String>, 22 22 + } 23 23 + 24 24 + async fn handle_discord_event(event: Event, _ctx: Arc<BotContext>) -> Result { 25 25 + match event { 26 26 + Event::MessageCreate(msg) => { 27 27 + let channel_id = msg.channel_id.to_string(); 28 28 + eprintln!("id: {channel_id}"); 29 29 + } 30 30 + Event::Ready(ev) => { 31 31 + eprintln!("Connected to gateway as {}", ev.user.name); 32 32 + } 33 33 + _ => {} 34 34 + } 35 35 + 36 36 + Ok(()) 37 37 + } 38 38 + 39 39 + #[tokio::main] 40 40 + async fn main() -> Result { 41 41 + // Config 42 42 + let token_file = std::env::var("TOKEN_FILE").context("Missing TOKEN_FILE env var")?; 43 43 + let reply_channels: HashSet<String> = HashSet::from_iter( 44 44 + std::env::var("REPLY_CHANNELS") 45 45 + .context("Missing REPLY_CHANNELS env var")? 46 46 + .split(",") 47 47 + .map(|s| s.trim().to_string()), 48 48 + ); 49 49 + let intents = Intents::GUILD_MESSAGES | Intents::MESSAGE_CONTENT; 50 50 + 51 51 + // Read token 52 52 + let token = std::fs::read_to_string(token_file).context("Failed to read bot token")?; 53 53 + let token = token.trim(); 54 54 + 55 55 + // Init 56 56 + let mut shard = Shard::new(ShardId::ONE, token.to_string(), intents); 57 57 + let http = HttpClient::new(token.to_string()); 58 58 + let cache = DefaultInMemoryCache::builder() 59 59 + .resource_types( 60 60 + ResourceType::MESSAGE 61 61 + | ResourceType::USER 62 62 + | ResourceType::CHANNEL 63 63 + | ResourceType::USER_CURRENT, 64 64 + ) 65 65 + .build(); 66 66 + 67 67 + let context = Arc::new(BotContext { 68 68 + http, 69 69 + reply_channels, 70 70 + }); 71 71 + 72 72 + // Event Loop 73 73 + while let Some(res) = shard.next_event(EventTypeFlags::all()).await { 74 74 + match res { 75 75 + Ok(event) => { 76 76 + cache.update(&event); 77 77 + tokio::spawn(handle_discord_event(event, Arc::clone(&context))); 78 78 + } 79 79 + Err(why) => { 80 80 + eprintln!("Failed to receive event: {why:?}"); 81 81 + } 82 82 + } 83 83 + } 84 84 + 85 85 + Ok(()) 3 86 }