first commit
This commit is contained in:
180
src/database.ts
Normal file
180
src/database.ts
Normal file
@@ -0,0 +1,180 @@
|
||||
import Database from 'better-sqlite3';
|
||||
import { MarkovChain } from './markov.js';
|
||||
|
||||
const DB_FILE = 'data/ulfbot.db';
|
||||
|
||||
// Cleanup settings
|
||||
const MAX_TRANSITIONS_PER_CHAT = 10000; // Max transitions per chat
|
||||
const MIN_TRANSITION_COUNT = 2; // Remove transitions seen only once
|
||||
const MIN_TRANSITIONS_TO_CLEANUP = 100; // Don't cleanup chats with fewer transitions
|
||||
|
||||
let db: Database.Database;
|
||||
|
||||
export interface ChatSettings {
|
||||
probability: number;
|
||||
}
|
||||
|
||||
export function initDatabase(): void {
|
||||
db = new Database(DB_FILE);
|
||||
|
||||
db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS chat_settings (
|
||||
chat_id INTEGER PRIMARY KEY,
|
||||
probability INTEGER DEFAULT 10
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS markov_transitions (
|
||||
chat_id INTEGER,
|
||||
key TEXT,
|
||||
next_word TEXT,
|
||||
count INTEGER,
|
||||
PRIMARY KEY (chat_id, key, next_word)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS markov_starts (
|
||||
chat_id INTEGER,
|
||||
key TEXT,
|
||||
count INTEGER,
|
||||
PRIMARY KEY (chat_id, key)
|
||||
);
|
||||
`);
|
||||
|
||||
// Create indexes for faster queries
|
||||
db.exec(`
|
||||
CREATE INDEX IF NOT EXISTS idx_transitions_chat ON markov_transitions(chat_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_starts_chat ON markov_starts(chat_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_transitions_count ON markov_transitions(count);
|
||||
`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Remove low-frequency transitions and limit total size per chat
|
||||
* Only cleans up chats with enough data (MIN_TRANSITIONS_TO_CLEANUP)
|
||||
*/
|
||||
export function cleanupDatabase(): void {
|
||||
const transaction = db.transaction(() => {
|
||||
// Only cleanup chats that have enough transitions
|
||||
// Chats with fewer transitions are still learning - don't delete their data
|
||||
db.prepare(`
|
||||
DELETE FROM markov_transitions
|
||||
WHERE count < ? AND chat_id IN (
|
||||
SELECT chat_id FROM markov_transitions
|
||||
GROUP BY chat_id HAVING COUNT(*) >= ?
|
||||
)
|
||||
`).run(MIN_TRANSITION_COUNT, MIN_TRANSITIONS_TO_CLEANUP);
|
||||
|
||||
// Get chats that exceed the limit
|
||||
const chats = db.prepare(`
|
||||
SELECT chat_id, COUNT(*) as cnt FROM markov_transitions
|
||||
GROUP BY chat_id HAVING cnt > ?
|
||||
`).all(MAX_TRANSITIONS_PER_CHAT) as { chat_id: number; cnt: number }[];
|
||||
|
||||
for (const { chat_id } of chats) {
|
||||
// Keep only the most frequent transitions
|
||||
db.prepare(`
|
||||
DELETE FROM markov_transitions
|
||||
WHERE chat_id = ? AND (key, next_word) NOT IN (
|
||||
SELECT key, next_word FROM markov_transitions
|
||||
WHERE chat_id = ?
|
||||
ORDER BY count DESC
|
||||
LIMIT ?
|
||||
)
|
||||
`).run(chat_id, chat_id, MAX_TRANSITIONS_PER_CHAT);
|
||||
}
|
||||
|
||||
// Remove orphaned starts (no matching transitions)
|
||||
db.prepare(`
|
||||
DELETE FROM markov_starts
|
||||
WHERE key NOT IN (SELECT DISTINCT key FROM markov_transitions)
|
||||
`).run();
|
||||
});
|
||||
|
||||
transaction();
|
||||
|
||||
const stats = db.prepare(`
|
||||
SELECT
|
||||
(SELECT COUNT(*) FROM markov_transitions) as transitions,
|
||||
(SELECT COUNT(*) FROM markov_starts) as starts
|
||||
`).get() as { transitions: number; starts: number };
|
||||
|
||||
console.log(`DB cleanup done: ${stats.transitions} transitions, ${stats.starts} starts`);
|
||||
}
|
||||
|
||||
export function getSettings(chatId: number): ChatSettings {
|
||||
const row = db
|
||||
.prepare('SELECT probability FROM chat_settings WHERE chat_id = ?')
|
||||
.get(chatId);
|
||||
|
||||
return { probability: (row as { probability: number } | undefined)?.probability ?? 10 };
|
||||
}
|
||||
|
||||
export function setProbability(chatId: number, probability: number): void {
|
||||
db
|
||||
.prepare(`
|
||||
INSERT INTO chat_settings (chat_id, probability)
|
||||
VALUES (?, ?)
|
||||
ON CONFLICT(chat_id) DO UPDATE SET probability = excluded.probability
|
||||
`)
|
||||
.run(chatId, probability);
|
||||
}
|
||||
|
||||
export function loadChain(chatId: number): MarkovChain {
|
||||
const chain = new MarkovChain(2);
|
||||
|
||||
// Load transitions
|
||||
const transitions = db
|
||||
.prepare('SELECT key, next_word, count FROM markov_transitions WHERE chat_id = ?')
|
||||
.all(chatId) as { key: string; next_word: string; count: number }[];
|
||||
|
||||
for (const row of transitions) {
|
||||
chain.addTransition(row.key, row.next_word, row.count);
|
||||
}
|
||||
|
||||
// Load starts
|
||||
const starts = db
|
||||
.prepare('SELECT key, count FROM markov_starts WHERE chat_id = ?')
|
||||
.all(chatId) as { key: string; count: number }[];
|
||||
|
||||
for (const row of starts) {
|
||||
chain.addStart(row.key, row.count);
|
||||
}
|
||||
|
||||
return chain;
|
||||
}
|
||||
|
||||
export function saveChain(chatId: number, chain: MarkovChain): void {
|
||||
const data = chain.export();
|
||||
|
||||
// Use transaction for atomicity
|
||||
const saveTransaction = db.transaction(() => {
|
||||
// Clear existing data for this chat
|
||||
db.prepare('DELETE FROM markov_transitions WHERE chat_id = ?').run(chatId);
|
||||
db.prepare('DELETE FROM markov_starts WHERE chat_id = ?').run(chatId);
|
||||
|
||||
// Insert transitions
|
||||
const insertTransition = db.prepare(
|
||||
'INSERT INTO markov_transitions (chat_id, key, next_word, count) VALUES (?, ?, ?, ?)'
|
||||
);
|
||||
|
||||
for (const [key, words] of Object.entries(data.transitions)) {
|
||||
for (const [nextWord, count] of Object.entries(words as Record<string, number>)) {
|
||||
insertTransition.run(chatId, key, nextWord, count);
|
||||
}
|
||||
}
|
||||
|
||||
// Insert starts
|
||||
const insertStart = db.prepare(
|
||||
'INSERT INTO markov_starts (chat_id, key, count) VALUES (?, ?, ?)'
|
||||
);
|
||||
|
||||
for (const [key, count] of Object.entries(data.starts)) {
|
||||
insertStart.run(chatId, key, count);
|
||||
}
|
||||
});
|
||||
|
||||
saveTransaction();
|
||||
}
|
||||
|
||||
export function closeDatabase(): void {
|
||||
db.close();
|
||||
}
|
||||
173
src/index.ts
Normal file
173
src/index.ts
Normal file
@@ -0,0 +1,173 @@
|
||||
import 'dotenv/config';
|
||||
import { Telegraf, Context } from 'telegraf';
|
||||
import { MarkovChain } from './markov.js';
|
||||
import { initDatabase, closeDatabase, getSettings, setProbability, loadChain, saveChain, cleanupDatabase } from './database.js';
|
||||
|
||||
const bot = new Telegraf(process.env.BOT_TOKEN!);
|
||||
|
||||
// In-memory cache of chains
|
||||
const chains = new Map<number, MarkovChain>();
|
||||
|
||||
// Get or create chain for a chat (loads from DB if not cached)
|
||||
function getChain(chatId: number): MarkovChain {
|
||||
if (!chains.has(chatId)) {
|
||||
chains.set(chatId, loadChain(chatId));
|
||||
}
|
||||
return chains.get(chatId)!;
|
||||
}
|
||||
|
||||
// Check if user is admin in group
|
||||
async function isAdmin(ctx: Context, userId: number): Promise<boolean> {
|
||||
if (ctx.chat?.type === 'private') return true;
|
||||
|
||||
try {
|
||||
const admins = await ctx.getChatAdministrators();
|
||||
return admins.some((admin) => admin.user.id === userId);
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Get bot info to extract username
|
||||
let botUsername: string | null = null;
|
||||
|
||||
bot.use(async (ctx, next) => {
|
||||
if (!botUsername) {
|
||||
const me = await ctx.telegram.getMe();
|
||||
botUsername = me.username ?? null;
|
||||
}
|
||||
return next();
|
||||
});
|
||||
|
||||
bot.command('start', (ctx) => {
|
||||
ctx.reply(
|
||||
'Hallo! Ich lerne von Nachrichten und generiere neue Sätze.\n\n' +
|
||||
'Ich antworte wenn:\n' +
|
||||
'• Du mich per Reply markierst\n' +
|
||||
'• Du meinen @namen erwähnst\n\n' +
|
||||
'In Gruppen antworte ich zufällig mit eingestellter Wahrscheinlichkeit.\n\n' +
|
||||
'Ich lerne nur aus Nachrichten der jeweiligen Gruppe - jede Gruppe hat ihre eigene "Persönlichkeit".\n\n' +
|
||||
'/hilfe - Diese Hilfe anzeigen\n' +
|
||||
'/prob - Aktuelle Wahrscheinlichkeit anzeigen\n' +
|
||||
'/setprob <0-100> - Antwortwahrscheinlichkeit setzen (Admins only)'
|
||||
);
|
||||
});
|
||||
|
||||
bot.command('hilfe', (ctx) => {
|
||||
ctx.reply(
|
||||
'📚 Hilfe\n\n' +
|
||||
'Ich generiere Sätze basierend auf vorherigen Nachrichten in dieser Gruppe.\n\n' +
|
||||
'Befehle:\n' +
|
||||
'/start - Begrüßung anzeigen\n' +
|
||||
'/hilfe - Diese Hilfe\n' +
|
||||
'/prob - Antwortwahrscheinlichkeit anzeigen\n' +
|
||||
'/setprob <0-100> - Wahrscheinlichkeit setzen (Admins)\n' +
|
||||
'/cleanup - Datenbank bereinigen (Admins)\n\n' +
|
||||
'Trigger für Antwort:\n' +
|
||||
'• Reply auf eine meiner Nachrichten\n' +
|
||||
'• @Erwähnung meines Namens\n' +
|
||||
'• Zufällig (je nach Wahrscheinlichkeit)\n\n' +
|
||||
'Ich lerne nur aus Nachrichten dieser Gruppe.'
|
||||
);
|
||||
});
|
||||
|
||||
bot.command('prob', (ctx) => {
|
||||
const settings = getSettings(ctx.chat.id);
|
||||
ctx.reply(`Aktuelle Antwortwahrscheinlichkeit: ${settings.probability}%`);
|
||||
});
|
||||
|
||||
bot.command('setprob', async (ctx) => {
|
||||
if (!ctx.from || !(await isAdmin(ctx, ctx.from.id))) {
|
||||
ctx.reply('Nur Admins können die Wahrscheinlichkeit ändern.');
|
||||
return;
|
||||
}
|
||||
|
||||
const args = ctx.message?.text?.split(' ')[1];
|
||||
const prob = parseInt(args ?? '', 10);
|
||||
|
||||
if (isNaN(prob) || prob < 0 || prob > 100) {
|
||||
ctx.reply('Bitte eine Zahl zwischen 0 und 100 angeben.');
|
||||
return;
|
||||
}
|
||||
|
||||
setProbability(ctx.chat.id, prob);
|
||||
ctx.reply(`Wahrscheinlichkeit auf ${prob}% gesetzt.`);
|
||||
});
|
||||
|
||||
bot.command('cleanup', async (ctx) => {
|
||||
if (!ctx.from || !(await isAdmin(ctx, ctx.from.id))) {
|
||||
ctx.reply('Nur Admins können Cleanup ausführen.');
|
||||
return;
|
||||
}
|
||||
|
||||
cleanupDatabase();
|
||||
ctx.reply('Datenbank bereinigt: Seltene Einträge entfernt.');
|
||||
});
|
||||
|
||||
// Handle text messages - learn and potentially respond
|
||||
bot.on('text', async (ctx) => {
|
||||
if (!ctx.message || !ctx.from) return;
|
||||
|
||||
const text = ctx.message.text;
|
||||
|
||||
// Skip commands
|
||||
if (text.startsWith('/')) return;
|
||||
|
||||
// Learn from message
|
||||
const chain = getChain(ctx.chat.id);
|
||||
chain.learn(text);
|
||||
saveChain(ctx.chat.id, chain);
|
||||
|
||||
// Check if bot is mentioned or replied to
|
||||
const botId = (await ctx.telegram.getMe()).id;
|
||||
const isReplyToBot = ctx.message.reply_to_message?.from?.id === botId;
|
||||
const isMentioned = botUsername && text.toLowerCase().includes(`@${botUsername.toLowerCase()}`);
|
||||
|
||||
// Determine if should respond
|
||||
let shouldRespond = false;
|
||||
|
||||
if (isReplyToBot || isMentioned) {
|
||||
shouldRespond = true;
|
||||
} else if (ctx.chat.type !== 'private') {
|
||||
const settings = getSettings(ctx.chat.id);
|
||||
if (Math.random() * 100 < settings.probability) {
|
||||
shouldRespond = true;
|
||||
}
|
||||
} else {
|
||||
shouldRespond = true;
|
||||
}
|
||||
|
||||
if (shouldRespond && chain.hasLearned()) {
|
||||
const sentence = chain.generate();
|
||||
ctx.reply(sentence, { reply_parameters: { message_id: ctx.message.message_id } });
|
||||
}
|
||||
});
|
||||
|
||||
// Graceful shutdown
|
||||
process.once('SIGINT', () => {
|
||||
bot.stop('SIGINT');
|
||||
closeDatabase();
|
||||
});
|
||||
process.once('SIGTERM', () => {
|
||||
bot.stop('SIGTERM');
|
||||
closeDatabase();
|
||||
});
|
||||
|
||||
// Periodic cleanup (every 24 hours)
|
||||
const CLEANUP_INTERVAL_MS = 24 * 60 * 60 * 1000;
|
||||
|
||||
function startPeriodicCleanup(): void {
|
||||
setInterval(() => {
|
||||
console.log('Running periodic cleanup...');
|
||||
cleanupDatabase();
|
||||
// Clear in-memory cache to free memory
|
||||
chains.clear();
|
||||
}, CLEANUP_INTERVAL_MS);
|
||||
}
|
||||
|
||||
// Start bot
|
||||
initDatabase();
|
||||
cleanupDatabase();
|
||||
bot.launch();
|
||||
startPeriodicCleanup();
|
||||
console.log('Bot is running...');
|
||||
170
src/markov.ts
Normal file
170
src/markov.ts
Normal file
@@ -0,0 +1,170 @@
|
||||
/**
|
||||
* Markov chain for generating sentences from learned text.
|
||||
* Uses trigrams (word triplets) for better grammatical coherence.
|
||||
*/
|
||||
|
||||
interface ChainData {
|
||||
// Map of "word1 word2" -> possible next words with frequency
|
||||
transitions: Map<string, Map<string, number>>;
|
||||
// All sentence start patterns
|
||||
starts: Map<string, number>;
|
||||
}
|
||||
|
||||
export class MarkovChain {
|
||||
private order: number;
|
||||
private chain: ChainData;
|
||||
|
||||
constructor(order: number = 2) {
|
||||
this.order = order;
|
||||
this.chain = {
|
||||
transitions: new Map(),
|
||||
starts: new Map(),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Tokenize text into words, preserving sentence boundaries
|
||||
*/
|
||||
private tokenize(text: string): string[][] {
|
||||
const sentences: string[][] = [];
|
||||
|
||||
// Split into sentences (rough but works for most cases)
|
||||
const sentencePatterns = text.split(/[.!?]+/);
|
||||
|
||||
for (const sentence of sentencePatterns) {
|
||||
const words = sentence
|
||||
.trim()
|
||||
.toLowerCase()
|
||||
.replace(/[^\wäöüß\s]/g, '') // Keep German characters
|
||||
.split(/\s+/)
|
||||
.filter(w => w.length > 0);
|
||||
|
||||
if (words.length > 0) {
|
||||
sentences.push(words);
|
||||
}
|
||||
}
|
||||
|
||||
return sentences;
|
||||
}
|
||||
|
||||
/**
|
||||
* Learn from a text, adding it to the chain
|
||||
*/
|
||||
learn(text: string): void {
|
||||
const sentences = this.tokenize(text);
|
||||
|
||||
for (const words of sentences) {
|
||||
if (words.length < this.order + 1) continue;
|
||||
|
||||
// Mark sentence start
|
||||
const startKey = words.slice(0, this.order).join(' ');
|
||||
this.chain.starts.set(startKey, (this.chain.starts.get(startKey) || 0) + 1);
|
||||
|
||||
// Build transitions
|
||||
for (let i = 0; i < words.length - this.order; i++) {
|
||||
const key = words.slice(i, i + this.order).join(' ');
|
||||
const next = words[i + this.order];
|
||||
|
||||
if (!this.chain.transitions.has(key)) {
|
||||
this.chain.transitions.set(key, new Map());
|
||||
}
|
||||
|
||||
const transitions = this.chain.transitions.get(key)!;
|
||||
transitions.set(next, (transitions.get(next) || 0) + 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Pick a weighted random item from a map of {item: weight}
|
||||
*/
|
||||
private weightedRandom<T>(items: Map<T, number>): T | null {
|
||||
const entries = Array.from(items.entries());
|
||||
if (entries.length === 0) return null;
|
||||
|
||||
const total = entries.reduce((sum, [, w]) => sum + w, 0);
|
||||
let random = Math.random() * total;
|
||||
|
||||
for (const [item, weight] of entries) {
|
||||
random -= weight;
|
||||
if (random <= 0) return item;
|
||||
}
|
||||
|
||||
return entries[0]![0];
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a sentence from the learned chain
|
||||
*/
|
||||
generate(maxWords: number = 20): string {
|
||||
if (this.chain.starts.size === 0) {
|
||||
return 'Ich brauche erst mehr Nachrichten zum Lernen.';
|
||||
}
|
||||
|
||||
// Pick random start
|
||||
const startKey = this.weightedRandom(this.chain.starts);
|
||||
if (!startKey) return '...';
|
||||
|
||||
const words: string[] = startKey.split(' ');
|
||||
let currentKey = startKey;
|
||||
|
||||
while (words.length < maxWords) {
|
||||
const transitions = this.chain.transitions.get(currentKey);
|
||||
|
||||
if (!transitions || transitions.size === 0) break;
|
||||
|
||||
const next = this.weightedRandom(transitions);
|
||||
if (!next) break;
|
||||
|
||||
words.push(next);
|
||||
|
||||
// Update key (slide window)
|
||||
const keyWords = words.slice(-this.order);
|
||||
currentKey = keyWords.join(' ');
|
||||
}
|
||||
|
||||
// Capitalize first letter
|
||||
const result = words.join(' ');
|
||||
return result.charAt(0).toUpperCase() + result.slice(1);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if chain has learned anything
|
||||
*/
|
||||
hasLearned(): boolean {
|
||||
return this.chain.starts.size > 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a single transition (for loading from DB)
|
||||
*/
|
||||
addTransition(key: string, nextWord: string, count: number): void {
|
||||
if (!this.chain.transitions.has(key)) {
|
||||
this.chain.transitions.set(key, new Map());
|
||||
}
|
||||
this.chain.transitions.get(key)!.set(nextWord, count);
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a single start (for loading from DB)
|
||||
*/
|
||||
addStart(key: string, count: number): void {
|
||||
this.chain.starts.set(key, count);
|
||||
}
|
||||
|
||||
/**
|
||||
* Export chain data for persistence
|
||||
*/
|
||||
export(): { transitions: Record<string, Record<string, number>>; starts: Record<string, number> } {
|
||||
const transitions: Record<string, Record<string, number>> = {};
|
||||
|
||||
for (const [key, words] of this.chain.transitions) {
|
||||
transitions[key] = Object.fromEntries(words);
|
||||
}
|
||||
|
||||
return {
|
||||
transitions,
|
||||
starts: Object.fromEntries(this.chain.starts),
|
||||
};
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user