diff --git a/src/brain.rs b/src/brain.rs index 4cbf63c..d7a1598 100644 --- a/src/brain.rs +++ b/src/brain.rs @@ -153,6 +153,15 @@ impl Brain { } } + pub fn forget_edge(&mut self, from: &str, to: &str) -> bool { + if let Some(edges) = self.0.get_mut(&Self::normalize_token(from)) { + edges.forget(&Self::normalize_token(to)); + true + } else { + false + } + } + pub fn merge_from(&mut self, other: Self) { for (k, v) in other.0.into_iter() { if let Some(edges) = self.0.get_mut(&k) { @@ -382,6 +391,37 @@ mod tests { assert_eq!(edges.1, 1); } + #[test] + fn forget_edge() { + let mut brain = Brain::default(); + + brain.ingest("hello world"); + brain.ingest("hello evil"); + brain.ingest("evil bad"); + + let exists = brain.forget_edge("hello", "evil"); + + assert!(exists, "hello -> evil did not exist"); + + assert!( + brain.0.contains_key(&Some(String::from("evil"))), + "Edges don't exist for evil" + ); + let edges = brain + .0 + .get(&Some(String::from("hello"))) + .expect("No weights for hello"); + assert!( + !edges.0.contains_key(&Some(String::from("evil"))), + "Edges for hello still has evil" + ); + assert!( + edges.0.contains_key(&Some(String::from("world"))), + "Edges for hello does not have world" + ); + assert_eq!(edges.1, 1); + } + #[test] fn none_on_empty() { let mut brain = Brain::default(); diff --git a/src/cmd/forget_edge.rs b/src/cmd/forget_edge.rs new file mode 100644 index 0000000..11c870b --- /dev/null +++ b/src/cmd/forget_edge.rs @@ -0,0 +1,60 @@ +use std::sync::{Arc, atomic::Ordering}; + +use twilight_interactions::command::{CommandModel, CreateCommand}; +use twilight_model::application::interaction::{Interaction, application_command::CommandData}; + +use crate::{ + BotContext, cmd::DEFER_INTER_RESP_EPHEMERAL, prelude::*, require_owner, status::update_status, +}; + +#[derive(CommandModel, CreateCommand)] +#[command( + name = "forget-edge", + desc = "Erase a specific edge in the graph. THIS ACTION IS IRREVERSIBLE!" +)] +pub struct ForgetEdgeCommand { + /// From token + from: String, + /// To token + to: String, +} + +impl ForgetEdgeCommand { + pub async fn handle(inter: Interaction, data: CommandData, ctx: Arc) -> Result { + let client = ctx.http.interaction(ctx.app_id); + + require_owner!(inter, ctx, client); + + let Self { from, to } = + Self::from_interaction(data.into()).context("Failed to parse command data")?; + + client + .create_response(inter.id, &inter.token, &DEFER_INTER_RESP_EPHEMERAL) + .await + .context("Failed to defer")?; + + let existed = { + let mut brain = ctx.brain_handle.write().await; + let existed = brain.forget_edge(&from, &to); + if existed { + ctx.pending_save.store(true, Ordering::Relaxed); + update_status(&brain, &ctx.shard_sender).context("Failed to update status")?; + } + existed + }; + + let msg = if existed { + "Edge forgotten" + } else { + "That edge does not seem to exist" + }; + + client + .update_response(&inter.token) + .content(Some(msg)) + .await + .context("Failed to send brain")?; + + Ok(()) + } +} diff --git a/src/cmd/mod.rs b/src/cmd/mod.rs index 4b7c0f1..3ca1c0f 100644 --- a/src/cmd/mod.rs +++ b/src/cmd/mod.rs @@ -1,5 +1,6 @@ mod dump_chain; mod forget; +mod forget_edge; mod load_chain; mod markov; mod weights; @@ -18,6 +19,7 @@ use crate::{BotContext, prelude::*}; use dump_chain::DumpChainCommand; use forget::ForgetCommand; +use forget_edge::ForgetEdgeCommand; use load_chain::LoadChainCommand; use markov::MarkovCommand; use weights::WeightsCommand; @@ -69,6 +71,7 @@ pub async fn register_all_commands(ctx: Arc) -> Result { LoadChainCommand::create_command().into(), MarkovCommand::create_command().into(), ForgetCommand::create_command().into(), + ForgetEdgeCommand::create_command().into(), ]; let client = ctx.http.interaction(ctx.app_id); @@ -92,6 +95,7 @@ pub async fn handle_app_command( "load_chain" => LoadChainCommand::handle(inter, data, ctx).await, "markov" => MarkovCommand::handle(inter, data, ctx).await, "forget" => ForgetCommand::handle(inter, data, ctx).await, + "forget-edge" => ForgetEdgeCommand::handle(inter, data, ctx).await, other => { warn!("Unknown command send: {other}"); Ok(())