awesome new feature

Author: Tricked-dev

Cargo.lock

@@ -41,10 +41,11 @@ color-eyre = "0.6.3"
 vesper = "0.13.0"
 once_cell = { version = "1.19.0", features = ["parking_lot"] }
 serde_rusqlite = "0.39.0"
-rig-core = "0.13.0"
+rig = { package = "rig-core", version = "0.13.0" }
 thiserror = "2.0.12"
 r2d2_sqlite = "0.30.0"
 r2d2 = "0.8.10"
 wb_sqlite = "0.2.1"
 rusqlite = { version = "0.36.0", features = ["bundled"] }
 num-format = "0.4.4"
+fasteval = "0.2.4"

Cargo.toml

@@ -81,3 +81,14 @@ pub struct Memory {
     #[sql(constraint = "NOT NULL")]
     pub key: String,
 }
+
+#[derive(CreateTableSql, InsertSync, UpdateSync, Serialize, Deserialize, Debug, PartialEq, Clone)]
+pub struct MathQuestion {
+    #[sql(constraint = "PRIMARY KEY AUTOINCREMENT")]
+    #[sql(typ = "INTEGER")]
+    pub id: i32,
+    #[sql(constraint = "NOT NULL")]
+    pub question: String,
+    #[sql(constraint = "NOT NULL")]
+    pub answer: f64,
+}

src/database.rs

@@ -41,6 +41,7 @@ mod commands;
 mod config;
 mod database;
 mod event_handler;
+mod math_test;
 mod message_handler;
 mod structs;
 pub mod utils;
@@ -82,6 +83,7 @@ async fn main() -> color_eyre::Result<()> {
     }
 
     rusqlite.execute(database::Memory::CREATE_TABLE_SQL, [])?;
+    rusqlite.execute(database::MathQuestion::CREATE_TABLE_SQL, [])?;
 
     let config = Arc::new(cfg);
 

src/main.rs

@@ -0,0 +1,142 @@
+use crate::database::MathQuestion;
+use color_eyre::Result;
+use rand::Rng;
+use rig::{completion::Prompt, providers, client::CompletionClient};
+use r2d2::Pool;
+use r2d2_sqlite::SqliteConnectionManager;
+use serde_rusqlite::from_rows;
+
+const MATH_PROMPTS: [&str; 3] = [
+    "Generate a simple mental math expression that can be solved in your head. Use basic operations like addition, subtraction, multiplication, or division with small numbers (prefer numbers under 20, maximum 100). Output ONLY the mathematical expression, nothing else. Examples: {ex1}, {ex2}, {ex3}, {ex4}",
+    "Create an easy math problem with 2-4 numbers that someone can calculate mentally. Use addition, subtraction, multiplication, or simple division. Keep numbers small and friendly (single or double digits preferred). Output ONLY the expression. Examples: {ex1}, {ex2}, {ex3}, {ex4}",
+    "Write a simple arithmetic calculation using small, friendly numbers that can be computed without a calculator. Stick to basic operations (+, -, *, /). Output ONLY the math expression. Examples: {ex1}, {ex2}, {ex3}, {ex4}"
+];
+
+pub struct MathTest {
+    pub question: String,
+    pub answer: f64,
+}
+
+impl MathTest {
+    pub async fn generate(
+        openai_api_key: &str,
+        db: &Pool<SqliteConnectionManager>,
+        rng: &mut impl Rng,
+    ) -> Result<Self> {
+        let client = providers::openai::Client::new(openai_api_key);
+
+        // Retry loop to avoid recursion
+        for attempt in 0..5 {
+            // Generate random example questions to show the AI
+            // Example 1: Addition
+            let ex1_a = rng.gen_range(5..50);
+            let ex1_b = rng.gen_range(5..50);
+            let ex1 = format!("{} + {}", ex1_a, ex1_b);
+
+            // Example 2: Multiplication
+            let ex2_a = rng.gen_range(3..15);
+            let ex2_b = rng.gen_range(3..15);
+            let ex2 = format!("{} * {}", ex2_a, ex2_b);
+
+            // Example 3: Division or Subtraction
+            let ex3 = if rng.gen_bool(0.5) {
+                let divisor = rng.gen_range(2..13);
+                let result = rng.gen_range(5..20);
+                format!("{} / {}", divisor * result, divisor)
+            } else {
+                let ex3_a = rng.gen_range(30..100);
+                let ex3_b = rng.gen_range(5..30);
+                format!("{} - {}", ex3_a, ex3_b)
+            };
+
+            // Example 4: Multi-operation or simple operation
+            let ex4 = if rng.gen_bool(0.3) {
+                // Multi-operation
+                let a = rng.gen_range(3..20);
+                let b = rng.gen_range(3..20);
+                let c = rng.gen_range(3..20);
+                match rng.gen_range(0..3) {
+                    0 => format!("{} + {} + {}", a, b, c),
+                    1 => format!("{} - {} + {}", a + b + c, b, c),
+                    _ => format!("{} + {} - {}", a, b, c),
+                }
+            } else {
+                // Simple division
+                let divisor = rng.gen_range(2..11);
+                let result = rng.gen_range(5..15);
+                format!("{} / {}", divisor * result, divisor)
+            };
+
+            // Select a random prompt template and fill in the example questions
+            let prompt_template = MATH_PROMPTS[rng.gen_range(0..MATH_PROMPTS.len())];
+            let prompt = prompt_template
+                .replace("{ex1}", &ex1)
+                .replace("{ex2}", &ex2)
+                .replace("{ex3}", &ex3)
+                .replace("{ex4}", &ex4);
+
+            // Generate question using AI
+            let agent = client
+                .agent("gpt-4o")
+                .preamble(
+                    "You are a math expression generator for mental math challenges. \
+                    Generate simple arithmetic expressions that can be solved mentally. \
+                    Output ONLY the mathematical expression with numbers and operators, no explanations, no greetings, no additional text."
+                )
+                .temperature(0.9)
+                .max_tokens(50)
+                .build();
+
+            let question = match agent.prompt(&prompt).await {
+                Ok(q) => q.trim().to_string(),
+                Err(e) => {
+                    tracing::error!("AI generation failed on attempt {}: {:?}", attempt + 1, e);
+                    continue;
+                }
+            };
+
+            // Validate the expression with fasteval
+            let mut ns = fasteval::EmptyNamespace;
+            let answer = match fasteval::ez_eval(&question, &mut ns) {
+                Ok(result) => result,
+                Err(e) => {
+                    tracing::error!("Invalid math expression generated '{}': {:?}", question, e);
+                    continue;
+                }
+            };
+
+            // Check if this question already exists in the database
+            let conn = db.get()?;
+            let mut stmt = conn.prepare("SELECT * FROM math_question WHERE question = ?")?;
+            let existing: Result<Vec<MathQuestion>, _> = from_rows(stmt.query([&question])?).collect();
+
+            // If question exists, try again
+            if existing.is_ok() && !existing.as_ref().unwrap().is_empty() {
+                tracing::debug!("Question '{}' already exists, retrying", question);
+                continue;
+            }
+
+            // Store in database - use manual INSERT to let SQLite handle autoincrement
+            conn.execute(
+                "INSERT INTO math_question (question, answer) VALUES (?, ?)",
+                rusqlite::params![&question, &answer],
+            )?;
+
+            return Ok(MathTest { question, answer });
+        }
+
+        // If all attempts failed, return error
+        Err(color_eyre::eyre::eyre!("Failed to generate valid math question after 5 attempts"))
+    }
+
+    pub fn validate_answer(&self, user_answer: &str) -> bool {
+        // Parse user answer
+        let user_answer = match user_answer.trim().parse::<f64>() {
+            Ok(v) => v,
+            Err(_) => return false,
+        };
+
+        // Check if answer is within 0.1 tolerance (1 decimal place)
+        (self.answer - user_answer).abs() <= 0.1
+    }
+}

src/math_test.rs

@@ -6,14 +6,16 @@ use rand::{
 use serde_rusqlite::from_row;
 use std::{sync::Arc, time::Instant};
 use tokio::sync::MutexGuard;
+use tokio::time::Instant as TokioInstant;
 use twilight_http::Client as HttpClient;
 use twilight_model::{gateway::payload::incoming::MessageCreate, id::Id};
 use vesper::twilight_exports::UserMarker;
 
 use crate::{
     ai_message,
     database::User,
-    structs::{Command, List, State},
+    math_test::MathTest,
+    structs::{Command, List, PendingMathTest, State},
     utils::levels::xp_required_for_level,
     zalgos::zalgify_text,
     RESPONDERS,
@@ -33,6 +35,146 @@ pub async fn handle_message(
         }
     }
 
+    // Check if user has a pending math test
+    if let Some(pending_test) = locked_state.pending_math_tests.get(&msg.author.id.get()) {
+        let elapsed = pending_test.started_at.elapsed();
+
+        // Check if 30 seconds have passed
+        if elapsed.as_secs() > 30 {
+            // Failed - timeout
+            locked_state.pending_math_tests.remove(&msg.author.id.get());
+
+            // Timeout the user for 1 minute
+            let timeout_until = twilight_model::util::Timestamp::from_secs(
+                std::time::SystemTime::now()
+                    .duration_since(std::time::UNIX_EPOCH)
+                    .unwrap()
+                    .as_secs() as i64
+                    + 60,
+            )
+            .unwrap();
+
+            if let Some(guild_id) = msg.guild_id {
+                match http
+                    .update_guild_member(guild_id, msg.author.id)
+                    .communication_disabled_until(Some(timeout_until))
+                {
+                    Ok(req) => {
+                        if let Err(e) = req.exec().await {
+                            tracing::error!("Failed to execute timeout: {:?}", e);
+                        }
+                    }
+                    Err(e) => {
+                        tracing::error!("Failed to timeout user: {:?}", e);
+                    }
+                }
+            }
+
+            return Ok(Command::text(format!(
+                "<@{}> Time's up! You took too long to answer. You've been timed out for 1 minute.",
+                msg.author.id.get()
+            ))
+            .reply());
+        }
+
+        // Check if answer is correct
+        let user_answer = msg.content.trim();
+        if (MathTest {
+            question: pending_test.question.clone(),
+            answer: pending_test.answer,
+        })
+        .validate_answer(user_answer)
+        {
+            locked_state.pending_math_tests.remove(&msg.author.id.get());
+
+            // Award 50x normal XP (normal is 5-20, so this is 250-1000)
+            let bonus_xp = locked_state.rng.gen_range(250..1000);
+
+            // Update user XP
+            let db = locked_state.db.get()?;
+            let mut statement = db.prepare("SELECT * FROM user WHERE id = ?").unwrap();
+            if let Ok(mut user) = statement.query_one([msg.author.id.get().to_string()], |row| {
+                from_row::<User>(row).map_err(|_| rusqlite::Error::QueryReturnedNoRows)
+            }) {
+                let level = user.level;
+                let xp_required = xp_required_for_level(level);
+                let new_xp = user.xp + bonus_xp;
+                user.name = msg.author.name.clone();
+
+                if new_xp >= xp_required {
+                    let new_level = level + 1;
+                    user.level = new_level;
+                    user.xp = new_xp - xp_required;
+                    user.update_sync(&db)?;
+
+                    return Ok(Command::text(format!(
+                        "<@{}> Correct! Well done. You earned {} XP and leveled up to level {}!",
+                        msg.author.id.get(),
+                        bonus_xp,
+                        new_level
+                    ))
+                    .reply());
+                } else {
+                    user.xp = new_xp;
+                    user.update_sync(&db)?;
+                }
+            }
+
+            return Ok(Command::text(format!(
+                "<@{}> Correct! Well done. You earned {} XP!",
+                msg.author.id.get(),
+                bonus_xp
+            ))
+            .reply());
+        } else {
+            // Wrong answer - don't remove, they can keep trying until timeout
+            return Ok(Command::text(format!(
+                "<@{}> Wrong answer! Try again. (Time remaining: {} seconds)",
+                msg.author.id.get(),
+                30 - elapsed.as_secs()
+            ))
+            .reply());
+        }
+    }
+
+    // Random 1/100 chance to trigger math test
+    let should_trigger_math = locked_state.config.openai_api_key.is_some()
+        && locked_state.rng.gen_range(0..1) == 0
+        && !locked_state.pending_math_tests.contains_key(&msg.author.id.get());
+
+    if should_trigger_math {
+        let api_key = locked_state.config.openai_api_key.clone().unwrap();
+        let db_clone = locked_state.db.clone();
+
+        // Use a new RNG for the async operation
+        let mut new_rng = rand::thread_rng();
+
+        match MathTest::generate(&api_key, &db_clone, &mut new_rng).await {
+            Ok(test) => {
+                let pending = PendingMathTest {
+                    user_id: msg.author.id.get(),
+                    channel_id: msg.channel_id.get(),
+                    question: test.question.clone(),
+                    answer: test.answer,
+                    started_at: TokioInstant::now(),
+                };
+
+                locked_state.pending_math_tests.insert(msg.author.id.get(), pending);
+
+                return Ok(Command::text(format!(
+                    "<@{}> **MATH TEST TIME!** Solve this in 30 seconds:\n`{}`\n(Answer to 1 decimal place)",
+                    msg.author.id.get(),
+                    test.question
+                ))
+                .reply()
+                .mention());
+            }
+            Err(e) => {
+                tracing::error!("Failed to generate math test: {:?}", e);
+            }
+        }
+    }
+
     let user = {
         let db = locked_state.db.get()?;
         let mut statement = db.prepare("SELECT * FROM user WHERE id = ?").unwrap();