ewenlau/tg-dev-srv-bot
Public Access
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 8 Wiki Activity

add some new commands
Run cargo test / Run tests (push) Successful in 4m7s
Details
Build Docker Package / build (push) Successful in 6m47s
Details
Build Docker Package / build (release) Successful in 7m16s
Details

Browse Source
This commit is contained in:
Elias Wendland
2026-07-05 19:49:00 +02:00
parent afa64fd55b
commit b1116846aa
10 changed files with 543 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.toml
+3
View File
@@ -10,18 +10,21 @@ repository = "https://git.ewenlau.net/ewenlau/tg-dev-srv-bot"
homepage = "https://git.ewenlau.net/ewenlau/tg-dev-srv-bot"
[dependencies]
chrono = "0.4.45"
dotenv = "0.15.0"
exmex = "0.21.0"
f128 = "0.2.9"
f256 = "0.11.2"
num-bigint = "0.5.1"
num-traits = "0.2.19"
ollama-rs = { version = "0.3.5", features = ["stream"] }
poise = "0.6.2"
regex = "1.12.4"
rug = "1.30.0"
serenity = "0.12.5"
sqlx = { version = "0.9.0", features = ["postgres", "runtime-tokio", "macros"] }
tokio = { version = "1.52.3", features = ["full"] }
tokio-stream = "0.1.18"
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
exmex_test
BIN
View File
Binary file not shown.
src/commands/ai.rs
+295
View File
@@ -0,0 +1,295 @@
use crate::{Context, Error};
use tracing::{debug, error, trace};
use ollama_rs::Ollama;
use std::env;
use std::time::Duration;
use ollama_rs::models::create::CreateModelRequest;
use ollama_rs::generation::completion::request::GenerationRequest;
use chrono::{DateTime, Utc};
use tokio_stream::StreamExt;
use poise::CreateReply;
/// Prompt the super advanced EwiAI
#[poise::command(slash_command, prefix_command)]
pub async fn prompt(
ctx: Context<'_>,
#[description = "The prompt to send to EwiAI"]
prompt: String,
#[description = "Include recent messages (max 10)"]
include_messages: Option<u8>,
) -> Result<(), Error> {
debug!("{} has requested to prompt EwiAI with '{}'", ctx.author().name, prompt);
let host = match env::var("TG_BOT_OLLAMA_HOST") {
Ok(h) => h,
Err(_) => {
ctx.say("Error: Expected an ollama url in the environment (`TG_BOT_OLLAMA_HOST`).").await?;
return Ok(());
}
};
let formatted_host = if !host.starts_with("http://") && !host.starts_with("https://") {
format!("http://{}", host)
} else {
host
};
let ollama = Ollama::builder()
.host(&formatted_host)
.port(11434)
.build();
let model_list = match ollama.list_local_models().await {
Ok(m) => m,
Err(e) => {
ctx.say(format!("Error: Failed to connect to Ollama: {}", e)).await?;
return Ok(());
}
};
let mut needs_create = true;
let yesterday = Utc::now() - chrono::Duration::days(1);
for model in &model_list {
if model.name.starts_with("EwiAI") {
if let Ok(modified) = DateTime::parse_from_rfc3339(&model.modified_at) {
if modified.with_timezone(&Utc) > yesterday {
needs_create = false;
}
}
}
}
if needs_create {
let system_prompt = system_prompt();
let _ = ctx.say("EwiAI model is missing or out of date. Creating/updating model (this may take a bit)...").await?;
if let Err(e) = ollama.create_model(CreateModelRequest::new("EwiAI".into())
.system(system_prompt.into())
.from_model("gemma4:e2b-it-qat".into())).await {
ctx.say(format!("Error creating model: {}", e)).await?;
return Ok(());
}
}
let mut final_prompt = String::new();
if let Some(mut limit) = include_messages {
if limit > 10 {
limit = 10;
}
match ctx.channel_id().messages(ctx.http(), poise::serenity_prelude::GetMessages::new().limit(limit)).await {
Ok(msgs) => {
final_prompt.push_str("Recent channel messages:\n");
for msg in msgs.iter().rev() {
let mut content = msg.content.clone();
if content.len() > 100 {
content.truncate(100);
content.push_str("...");
}
final_prompt.push_str(&format!("{} (ID: {}): {}\n", msg.author.name, msg.author.id, content));
}
final_prompt.push_str("\n");
}
Err(e) => {
error!("Failed to get messages: {:?}", e);
}
}
}
final_prompt.push_str(&format!("The user talking to you is {} (ID: {}).\nThat user is telling you the following: {}", ctx.author().name, ctx.author().id, prompt));
trace!("Final prompt: {}", final_prompt);
let reply = ctx.say("Thinking...\n-# The hardware this thing runs is really slow, expect a long wait").await?;
let mut stream = match ollama.generate_stream(GenerationRequest::new("EwiAI".into(), final_prompt).system(system_prompt())).await {
Ok(s) => s,
Err(e) => {
reply.edit(ctx, CreateReply::default().content(format!("Error during generation: {}", e))).await?;
return Ok(());
}
};
let mut response_text = String::new();
let mut last_update = std::time::Instant::now();
let mut final_response = None;
while let Some(res) = stream.next().await {
match res {
Ok(chunks) => {
for chunk in chunks {
response_text.push_str(&chunk.response);
if chunk.done {
final_response = Some(chunk);
}
}
if last_update.elapsed() >= Duration::from_secs(1) && !response_text.is_empty() {
let _ = reply.edit(ctx, CreateReply::default().content(&response_text)).await;
last_update = std::time::Instant::now();
}
}
Err(e) => {
let _ = reply.edit(ctx, CreateReply::default().content(format!("{} [Stream Error: {}]", response_text, e))).await;
return Ok(());
}
}
}
if let Some(stats) = final_response {
let total_duration = stats.total_duration.unwrap_or(0) as f64 / 1_000_000_000.0;
let eval_count = stats.eval_count.unwrap_or(0);
let eval_duration = stats.eval_duration.unwrap_or(0) as f64 / 1_000_000_000.0;
let tokens_per_sec = if eval_duration > 0.0 { eval_count as f64 / eval_duration } else { 0.0 };
let stats_text = format!(
"\n\n*Generated in {:.2}s ({:.2} tok/s)*",
total_duration, tokens_per_sec
);
response_text.push_str(&stats_text);
}
let _ = reply.edit(ctx, CreateReply::default().content(&response_text)).await;
Ok(())
}
/// Ask EwiAI to answer to recent messages
#[poise::command(slash_command, prefix_command)]
pub async fn answer(
ctx: Context<'_>,
#[description = "How many messages to include (max 10)"]
messages_count: Option<u8>,
) -> Result<(), Error> {
debug!("{} has requested EwiAI to answer", ctx.author().name);
let host = match env::var("TG_BOT_OLLAMA_HOST") {
Ok(h) => h,
Err(_) => {
ctx.say("Error: Expected an ollama url in the environment (`TG_BOT_OLLAMA_HOST`).").await?;
return Ok(());
}
};
let formatted_host = if !host.starts_with("http://") && !host.starts_with("https://") {
format!("http://{}", host)
} else {
host
};
let ollama = Ollama::builder()
.host(&formatted_host)
.port(11434)
.build();
let model_list = match ollama.list_local_models().await {
Ok(m) => m,
Err(e) => {
ctx.say(format!("Error: Failed to connect to Ollama: {}", e)).await?;
return Ok(());
}
};
let mut needs_create = true;
let yesterday = Utc::now() - chrono::Duration::days(1);
for model in &model_list {
if model.name.starts_with("EwiAI") {
if let Ok(modified) = DateTime::parse_from_rfc3339(&model.modified_at) {
if modified.with_timezone(&Utc) > yesterday {
needs_create = false;
}
}
}
}
if needs_create {
let system_prompt = system_prompt();
let _ = ctx.say("EwiAI model is missing or out of date. Creating/updating model (this may take a bit)...").await?;
if let Err(e) = ollama.create_model(CreateModelRequest::new("EwiAI".into())
.system(system_prompt.into())
.from_model("gemma4:e2b-it-qat".into())).await {
ctx.say(format!("Error creating model: {}", e)).await?;
return Ok(());
}
}
let mut final_prompt = String::new();
let mut limit = messages_count.unwrap_or(5);
if limit > 10 {
limit = 10;
}
match ctx.channel_id().messages(ctx.http(), poise::serenity_prelude::GetMessages::new().limit(limit)).await {
Ok(msgs) => {
final_prompt.push_str("Recent channel messages:\n");
for msg in msgs.iter().rev() {
let mut content = msg.content.clone();
if content.len() > 100 {
content.truncate(100);
content.push_str("...");
}
final_prompt.push_str(&format!("{} (ID: {}): {}\n", msg.author.name, msg.author.id, content));
}
final_prompt.push_str("\n");
}
Err(e) => {
error!("Failed to get messages: {:?}", e);
}
}
final_prompt.push_str(&format!("Please respond to the messages above."));
trace!("Final prompt: {}", final_prompt);
let reply = ctx.say("Thinking...\n-# The hardware this thing runs is really slow, expect a long wait").await?;
let mut stream = match ollama.generate_stream(GenerationRequest::new("EwiAI".into(), final_prompt).system(system_prompt())).await {
Ok(s) => s,
Err(e) => {
reply.edit(ctx, CreateReply::default().content(format!("Error during generation: {}", e))).await?;
return Ok(());
}
};
let mut response_text = String::new();
let mut last_update = std::time::Instant::now();
let mut final_response = None;
while let Some(res) = stream.next().await {
match res {
Ok(chunks) => {
for chunk in chunks {
response_text.push_str(&chunk.response);
if chunk.done {
final_response = Some(chunk);
}
}
if last_update.elapsed() >= Duration::from_secs(1) && !response_text.is_empty() {
let _ = reply.edit(ctx, CreateReply::default().content(&response_text)).await;
last_update = std::time::Instant::now();
}
}
Err(e) => {
let _ = reply.edit(ctx, CreateReply::default().content(format!("{} [Stream Error: {}]", response_text, e))).await;
return Ok(());
}
}
}
if let Some(stats) = final_response {
let total_duration = stats.total_duration.unwrap_or(0) as f64 / 1_000_000_000.0;
let eval_count = stats.eval_count.unwrap_or(0);
let eval_duration = stats.eval_duration.unwrap_or(0) as f64 / 1_000_000_000.0;
let tokens_per_sec = if eval_duration > 0.0 { eval_count as f64 / eval_duration } else { 0.0 };
let stats_text = format!(
"\n\n*Generated in {:.2}s ({:.2} tok/s)*",
total_duration, tokens_per_sec
);
response_text.push_str(&stats_text);
}
let _ = reply.edit(ctx, CreateReply::default().content(&response_text)).await;
Ok(())
}
fn system_prompt() -> String {
format!("You are EwiAI, an AI in the TeenGovernment Development Server. Your developer is Ewi/ewenlau (discord user id: 713354021124964422, discord username: ewenlau). Be slightly unhelpful, in a sarcastic way, but still answer the question. The current date is {}", Utc::now().to_string())
}
src/commands/fastfetch.rs
+44
View File
@@ -0,0 +1,44 @@
use std::process::Command;
use crate::{Context, Error};
use tracing::{debug, warn, trace};
use crate::shared_functions::trace_message;
/// Run and output the contents of the fastfetch command on the machine running the bot
#[poise::command(slash_command, prefix_command)]
pub async fn fastfetch(ctx: Context<'_>) -> Result<(), Error> {
debug!("fastfetch command called by user {} in guild {}", ctx.author().id.get(), ctx.guild_id().unwrap().get());
let msg = "Processing...".to_string();
let res_msg = ctx.say(&msg).await?;
trace_message(msg.clone(), ctx.channel_id().to_string(), ctx.guild_id().unwrap().to_string()).await;
let cmd = Command::new("fastfetch").args(&["--raw", "true", "--logo", "none"]).output();
match cmd {
Err(e) => {
warn!("Error executing fastfetch: {}", e);
let msg = format!("Error: {}", e);
edit_response_message(&res_msg, ctx, msg.clone(), false).await?;
trace_message(msg, ctx.channel_id().to_string(), ctx.guild_id().unwrap().to_string()).await;
Ok(())
}
Ok(output) => {
let output_str = output.stdout.into_iter().map(|c| c as char).collect::<String>();
let msg = format!("```ansi\n{}\n```", output_str);
edit_response_message(&res_msg, ctx, msg.clone(), false).await?;
// Tracing the output would be a bad idea since it's really long and filled with ansi escape codes
// trace_message(msg, ctx.channel_id().to_string(), ctx.guild_id().unwrap().to_string()).await;
trace!("Saying fastfetch command output");
Ok(())
}
}
}
async fn edit_response_message<'a>(response_message: &poise::ReplyHandle<'a>, ctx: Context<'_>, content: String, silent: bool) -> Result<(), Error> {
if silent {
response_message.edit(ctx, poise::CreateReply::default().content(content).allowed_mentions(serenity::all::CreateAllowedMentions::new().empty_users())).await?;
} else {
response_message.edit(ctx, poise::CreateReply::default().content(content)).await?;
}
Ok(())
}
src/commands/find_prime.rs
+183
View File
@@ -0,0 +1,183 @@
use crate::{Context, Error};
use tracing::{debug, trace};
use crate::shared_functions::{trace_message, edit_response_message};
/// Find the nth prime number (single threaded)
#[derive(Debug, poise::ChoiceParameter, Clone, Copy)]
pub enum Method {
#[name = "Sieve"]
Sieve,
#[name = "Trial and Error"]
TrialAndError,
}
#[poise::command(slash_command, prefix_command)]
pub async fn find_prime(
ctx: Context<'_>,
#[description = "Which nth prime number to find"]
n: u64,
#[description = "Method to use"]
method: Option<Method>,
#[description = "Custom timeout in seconds (upper limit of 60s)"]
timeout: Option<u64>,
) -> Result<(), Error> {
let guild_id_str = ctx.guild_id().map(|g| g.get().to_string()).unwrap_or_else(|| "DM".to_string());
debug!("Find prime command called by user {} in guild {}", ctx.author().id.get(), guild_id_str);
trace!("Primes requested for {} with method {:?}", n, method);
let msg = "Calculating...".to_string();
let res_msg = ctx.say(msg.clone()).await?;
trace_message(msg, ctx.channel_id().to_string(), guild_id_str.clone()).await;
if n == 0 {
let msg = "Prime indices start at 1. Please provide a value greater than 0.".to_string();
edit_response_message(&res_msg, ctx, msg.clone(), false).await?;
trace_message(msg, ctx.channel_id().to_string(), guild_id_str).await;
return Ok(());
}
if timeout.is_some() && timeout.unwrap() > 60 {
let msg = "Timeout is too long. Maximum is 60 seconds.".to_string();
edit_response_message(&res_msg, ctx, msg.clone(), false).await?;
trace_message(msg, ctx.channel_id().to_string(), guild_id_str).await;
return Ok(());
}
let start_time = std::time::Instant::now();
let timeout_duration = std::time::Duration::from_secs(timeout.unwrap_or(20));
let (tx, rx) = tokio::sync::oneshot::channel();
tokio::task::spawn_blocking(move || {
let result = match method {
Some(Method::Sieve) => sieve(n),
Some(Method::TrialAndError) => super3(n),
None => super3(n),
};
let _ = tx.send(result);
});
let method_name = match method {
Some(Method::Sieve) => "Sieve",
Some(Method::TrialAndError) => "TrialAndError",
None => "TrialAndError",
};
let result = match tokio::time::timeout(timeout_duration, rx).await {
Ok(Ok(result)) => result,
Ok(Err(_)) => {
let msg = "Calculation thread panicked or was dropped unexpectedly.".to_string();
edit_response_message(&res_msg, ctx, msg.clone(), false).await?;
trace_message(msg, ctx.channel_id().to_string(), guild_id_str).await;
return Ok(());
}
Err(_) => {
let msg = format!("Calculation timed out after {:?}", timeout_duration);
edit_response_message(&res_msg, ctx, msg.clone(), false).await?;
trace_message(msg, ctx.channel_id().to_string(), guild_id_str).await;
return Ok(());
}
};
let duration = start_time.elapsed();
let msg = format!("The {}th prime number is {}\n-# Calculation time: {:?}, Method: {:?}", n, result, duration, method_name);
edit_response_message(&res_msg, ctx, msg.clone(), false).await?;
trace_message(msg, ctx.channel_id().to_string(), guild_id_str).await;
Ok(())
}
// My best trial and error algorithm, originally designed to find as many primes as possible in 1s
pub fn super3(limit: u64) -> u64 {
if limit == 0 {
return 0;
}
let mut primes = Vec::new();
primes.push(2);
if limit == 1 {
return 2;
}
primes.push(3);
if limit == 2 {
return 3;
}
let mut current_multiplier = 1;
while (primes.len() as u64) < limit {
for offset in [-1_i64, 1] {
let base = 6 * current_multiplier;
let current_number = if offset < 0 {
base - 1
} else {
base + 1
};
let sqrt = (current_number as f64).sqrt().round() as u64;
let mut is_prime = true;
for &p in primes.iter().filter(|&&x| x > 3).take_while(|&&x| x <= sqrt) {
if current_number % p == 0 {
is_prime = false;
break;
}
}
if is_prime {
primes.push(current_number);
if (primes.len() as u64) == limit {
return current_number;
}
}
}
current_multiplier += 1;
}
*primes.last().unwrap()
}
fn sieve(n: u64) -> u64 {
if n == 0 {
panic!("Prime indices start at 1");
}
if n == 1 { return 2; }
if n == 2 { return 3; }
if n == 3 { return 5; }
if n == 4 { return 7; }
if n == 5 { return 11; }
let n_f = n as f64;
let ln_n = n_f.ln();
let limit = if n >= 688383 {
// Extremely tight bound for large n
(n_f * (ln_n + n_f.ln().ln() - 1.0 + (n_f.ln().ln() - 0.9385) / ln_n)) as usize
} else if n >= 6 {
// Safe bound for medium n
(n_f * (ln_n + n_f.ln().ln())) as usize
} else {
12 // Fallback for tiny numbers, though handled above
};
// Sieve of Eratosthenes
let mut is_prime = vec![true; limit + 1];
is_prime[0] = false;
is_prime[1] = false;
let mut count = 0;
for p in 2..=limit {
if is_prime[p] {
count += 1;
if count == n {
return p as u64;
}
// Safe overflow check for p * p
if let Some(mut i) = p.checked_mul(p) {
while i <= limit {
is_prime[i] = false;
i += p;
}
}
}
}
unreachable!("If we reach here, the upper bound calculation failed.");
}
src/commands/manage_auto_role.rs
+1 -1
View File
@@ -1,7 +1,7 @@
use crate::{Context, Error};
use crate::shared_functions::get_all_admin_users;
use poise::serenity_prelude as sere;
use tracing::{error, info, debug, trace};
use tracing::{error, debug, trace, info, warn};
#[derive(Debug, poise::ChoiceParameter)]
pub enum AutoRoleOperationType {
src/commands/mod.rs
+7
View File
@@ -3,9 +3,16 @@ pub mod calc;
pub mod manage_admins;
pub mod manage_auto_role;
pub mod version;
pub mod fastfetch;
pub mod find_prime;
pub mod ai;
pub use ping::ping;
pub use calc::calc;
pub use manage_admins::manage_admins;
pub use manage_auto_role::manage_auto_role;
pub use version::version;
pub use fastfetch::fastfetch;
pub use find_prime::find_prime;
pub use ai::prompt;
pub use ai::answer;
src/main.rs
+1 -1
View File
@@ -87,7 +87,7 @@ async fn main() {
}
}
}),
commands: vec![commands::ping(), commands::calc(), commands::manage_admins(), commands::manage_auto_role(), commands::version()],
commands: vec![commands::ping(), commands::calc(), commands::manage_admins(), commands::manage_auto_role(), commands::version(), commands::fastfetch(), commands::find_prime(), commands::prompt(), commands::answer()],
..Default::default()
})
.setup(|ctx, _ready, framework| {
src/shared_functions.rs
+9
View File
@@ -57,4 +57,13 @@ pub async fn get_all_admin_users(ctx: Context<'_>) -> Result<Vec<sere::User>, Er
pub async fn trace_message(msg: String, channel: String, guild: String) {
trace!("Saying \"{}\" in channel {} in guild {}", msg, channel, guild);
}
pub async fn edit_response_message<'a>(response_message: &poise::ReplyHandle<'a>, ctx: Context<'_>, content: String, silent: bool) -> Result<(), Error> {
if silent {
response_message.edit(ctx, poise::CreateReply::default().content(content).allowed_mentions(serenity::all::CreateAllowedMentions::new().empty_users())).await?;
} else {
response_message.edit(ctx, poise::CreateReply::default().content(content)).await?;
}
Ok(())
}
wipe-db.sh
-30
View File
@@ -1,30 +0,0 @@
#!/bin/bash
# AI WRITTEN
set -e
# Ensure we run from the script's directory
cd "$(dirname "$0")"
# Load environment variables if they exist
if [ -f .env ]; then
set -a
source .env
set +a
fi
echo "Wiping database volume and restarting containers..."
docker compose -f docker-compose-dev.yml down -v
docker compose -f docker-compose-dev.yml up -d
echo "Waiting for PostgreSQL to be ready..."
until [ "$(docker inspect -f '{{.State.Health.Status}}' tg-dev-srv-bot-postgres 2>/dev/null)" = "healthy" ]; do
printf "."
sleep 1
done
echo ""
echo "PostgreSQL is healthy!"
echo "Applying database schema from sqlschema.txt..."
docker compose -f docker-compose-dev.yml exec -T postgres psql -U "${TG_BOT_POSTGRES_USER:-postgres}" -d "${TG_BOT_POSTGRES_DB:-tg_dev_srv_bot}" -c "$(cat sqlschema.txt)"
echo "Database reset and schema applied successfully!"