From 23ebd021f38b531039d8ea7678b7f26ac15123e2 Mon Sep 17 00:00:00 2001 From: Aidan Date: Mon, 30 Jun 2025 23:43:30 -0400 Subject: [PATCH] ai queue, better markdown parsing, refactor, better feedback --- README.md | 1 + src/commands/ai.ts | 319 ++++++++++++++++++++++-------------- src/commands/main.ts | 8 +- src/locales/english.json | 45 +++-- src/locales/portuguese.json | 45 +++-- src/plugins/verifyInput.ts | 28 ++-- 6 files changed, 273 insertions(+), 173 deletions(-) diff --git a/README.md b/README.md index e035285..ba6ecef 100644 --- a/README.md +++ b/README.md @@ -117,6 +117,7 @@ If you prefer to use Docker directly, you can use these instructions instead. - **handlerTimeout** (optional): How long handlers will wait before timing out. Set this high if using large AI models. - **flashModel** (optional): Which model will be used for /ask - **thinkingModel** (optional): Which model will be used for /think +- **updateEveryChars** (optional): The amount of chars until message update triggers (for streaming response) - **databaseUrl**: Database server configuration (see `.env.example`) - **botAdmins**: Put the ID of the people responsible for managing the bot. They can use some administrative + exclusive commands on any group. - **lastKey**: Last.fm API key, for use on `lastfm.js` functions, like see who is listening to what song and etc. diff --git a/src/commands/ai.ts b/src/commands/ai.ts index de69903..3781e9f 100644 --- a/src/commands/ai.ts +++ b/src/commands/ai.ts @@ -119,31 +119,17 @@ export const models: ModelInfo[] = [ } ]; -const enSystemPrompt = `You are a plaintext-only, helpful assistant called {botName}. -Current Date/Time (UTC): {date} - ---- - -Respond to the user's message: -{message}` - -const ptSystemPrompt = `Você é um assistente de texto puro e útil chamado {botName}. -Data/Hora atual (UTC): {date} - ---- - -Responda à mensagem do usuário: -{message}` - -async function usingSystemPrompt(ctx: TextContext, db: NodePgDatabase, botName: string): Promise { +async function usingSystemPrompt(ctx: TextContext, db: NodePgDatabase, botName: string, message: string): Promise { const user = await db.query.usersTable.findMany({ where: (fields, { eq }) => eq(fields.telegramId, String(ctx.from!.id)), limit: 1 }); if (user.length === 0) await ensureUserInDb(ctx, db); const userData = user[0]; const lang = userData?.languageCode || "en"; + const Strings = getStrings(lang); const utcDate = new Date().toISOString(); - const prompt = lang === "pt" - ? ptSystemPrompt.replace("{botName}", botName).replace("{date}", utcDate).replace("{message}", ctx.message.text) - : enSystemPrompt.replace("{botName}", botName).replace("{date}", utcDate).replace("{message}", ctx.message.text); + const prompt = Strings.ai.systemPrompt + .replace("{botName}", botName) + .replace("{date}", utcDate) + .replace("{message}", message); return prompt; } @@ -156,6 +142,51 @@ export function sanitizeForJson(text: string): string { .replace(/\t/g, '\\t') } +function sanitizeMarkdownForTelegram(text: string): string { + let sanitizedText = text; + + const replacements: string[] = []; + const addReplacement = (match: string): string => { + replacements.push(match); + return `___PLACEHOLDER_${replacements.length - 1}___`; + }; + + sanitizedText = sanitizedText.replace(/```([\s\S]*?)```/g, addReplacement); + sanitizedText = sanitizedText.replace(/`([^`]+)`/g, addReplacement); + sanitizedText = sanitizedText.replace(/\[([^\]]+)\]\(([^)]+)\)/g, addReplacement); + + const parts = sanitizedText.split(/(___PLACEHOLDER_\d+___)/g); + const processedParts = parts.map(part => { + if (part.match(/___PLACEHOLDER_\d+___/)) { + return part; + } else { + let processedPart = part; + processedPart = processedPart.replace(/^(#{1,6})\s+(.+)/gm, '*$2*'); + processedPart = processedPart.replace(/^(\s*)[-*]\s+/gm, '$1- '); + processedPart = processedPart.replace(/\*\*(.*?)\*\*/g, '*$1*'); + processedPart = processedPart.replace(/__(.*?)__/g, '*$1*'); + processedPart = processedPart.replace(/(^|\s)\*(?!\*)([^*]+?)\*(?!\*)/g, '$1_$2_'); + processedPart = processedPart.replace(/(^|\s)_(?!_)([^_]+?)_(?!_)/g, '$1_$2_'); + processedPart = processedPart.replace(/~~(.*?)~~/g, '~$1~'); + processedPart = processedPart.replace(/^\s*┃/gm, '>'); + processedPart = processedPart.replace(/^>\s?/gm, '> '); + + return processedPart; + } + }); + + sanitizedText = processedParts.join(''); + + sanitizedText = sanitizedText.replace(/___PLACEHOLDER_(\d+)___/g, (_, idx) => replacements[Number(idx)]); + + const codeBlockCount = (sanitizedText.match(/```/g) || []).length; + if (codeBlockCount % 2 !== 0) { + sanitizedText += '\n```'; + } + + return sanitizedText; +} + export async function preChecks() { const envs = [ "ollamaApi", @@ -232,7 +263,7 @@ function extractAxiosErrorMessage(error: unknown): string { } function escapeMarkdown(text: string): string { - return text.replace(/([*_])/g, '\\$1'); + return text.replace(/([_*\[\]()`>#\+\-=|{}.!~])/g, '\\$1'); } function containsUrls(text: string): boolean { @@ -244,10 +275,14 @@ async function getResponse(prompt: string, ctx: TextContext, replyGenerating: Me if (!ctx.chat) { return { success: false, - error: Strings.unexpectedErr.replace("{error}", "No chat found"), + error: Strings.unexpectedErr.replace("{error}", Strings.ai.noChatFound), }; } - const modelHeader = `🤖 *${model}* | 🌡️ *${aiTemperature}*\n\n`; + let status = Strings.ai.statusWaitingRender; + let modelHeader = Strings.ai.modelHeader + .replace("{model}", model) + .replace("{temperature}", aiTemperature) + .replace("{status}", status) + "\n\n"; const urlWarning = containsUrls(originalMessage) ? Strings.ai.urlWarning : ''; try { @@ -267,8 +302,9 @@ async function getResponse(prompt: string, ctx: TextContext, replyGenerating: Me ); let fullResponse = ""; let thoughts = ""; - let lastUpdate = Date.now(); + let lastUpdateCharCount = 0; let sentHeader = false; + let firstChunk = true; const stream: NodeJS.ReadableStream = aiResponse.data as any; for await (const chunk of stream) { const lines = chunk.toString().split('\n'); @@ -293,7 +329,6 @@ async function getResponse(prompt: string, ctx: TextContext, replyGenerating: Me logger.logThinking(ctx.chat.id, replyGenerating.message_id, false); } } - const now = Date.now(); if (ln.response) { if (model === thinking_model) { let patchedThoughts = ln.response; @@ -306,20 +341,51 @@ async function getResponse(prompt: string, ctx: TextContext, replyGenerating: Me } else { fullResponse += ln.response; } - if (now - lastUpdate >= 5000 || !sentHeader) { + if (firstChunk) { + status = Strings.ai.statusWaitingRender; + modelHeader = Strings.ai.modelHeader + .replace("{model}", model) + .replace("{temperature}", aiTemperature) + .replace("{status}", status) + "\n\n"; await rateLimiter.editMessageWithRetry( ctx, ctx.chat.id, replyGenerating.message_id, - modelHeader + urlWarning + fullResponse, + modelHeader + urlWarning + escapeMarkdown(fullResponse), { parse_mode: 'Markdown' } ); - lastUpdate = now; + lastUpdateCharCount = fullResponse.length; + sentHeader = true; + firstChunk = false; + continue; + } + const updateEveryChars = Number(process.env.updateEveryChars) || 100; + if (fullResponse.length - lastUpdateCharCount >= updateEveryChars || !sentHeader) { + await rateLimiter.editMessageWithRetry( + ctx, + ctx.chat.id, + replyGenerating.message_id, + modelHeader + urlWarning + escapeMarkdown(fullResponse), + { parse_mode: 'Markdown' } + ); + lastUpdateCharCount = fullResponse.length; sentHeader = true; } } } } + status = Strings.ai.statusRendering; + modelHeader = Strings.ai.modelHeader + .replace("{model}", model) + .replace("{temperature}", aiTemperature) + .replace("{status}", status) + "\n\n"; + await rateLimiter.editMessageWithRetry( + ctx, + ctx.chat.id, + replyGenerating.message_id, + modelHeader + urlWarning + escapeMarkdown(fullResponse), + { parse_mode: 'Markdown' } + ); return { success: true, response: fullResponse, @@ -360,7 +426,7 @@ async function getResponse(prompt: string, ctx: TextContext, replyGenerating: Me console.log(`[✨ AI] ${model} pulled successfully`); return { success: true, - response: `✅ Pulled ${escapeMarkdown(model)} successfully, please retry the command.`, + response: Strings.ai.pulled.replace("{model}", escapeMarkdown(model)), }; } } @@ -376,16 +442,18 @@ async function handleAiReply(ctx: TextContext, model: string, prompt: string, re const aiResponse = await getResponse(prompt, ctx, replyGenerating, model, aiTemperature, originalMessage); if (!aiResponse) return; if (!ctx.chat) return; - const modelHeader = `🤖 *${model}* | 🌡️ *${aiTemperature}*\n\n`; - - const urlWarning = containsUrls(originalMessage) ? Strings.ai.urlWarning : ''; - if (aiResponse.success && aiResponse.response) { + const status = Strings.ai.statusComplete; + const modelHeader = Strings.ai.modelHeader + .replace("{model}", model) + .replace("{temperature}", aiTemperature) + .replace("{status}", status) + "\n\n"; + const urlWarning = containsUrls(originalMessage) ? Strings.ai.urlWarning : ''; await rateLimiter.editMessageWithRetry( ctx, ctx.chat.id, replyGenerating.message_id, - modelHeader + urlWarning + aiResponse.response, + modelHeader + urlWarning + sanitizeMarkdownForTelegram(aiResponse.response), { parse_mode: 'Markdown' } ); return; @@ -425,109 +493,112 @@ export function getModelLabelByName(name: string): string { export default (bot: Telegraf, db: NodePgDatabase) => { const botName = bot.botInfo?.first_name && bot.botInfo?.last_name ? `${bot.botInfo.first_name} ${bot.botInfo.last_name}` : "Kowalski" - bot.command(["ask", "think"], spamwatchMiddleware, async (ctx) => { - if (!ctx.message || !('text' in ctx.message)) return - const isAsk = ctx.message.text.startsWith("/ask") - const model = isAsk ? flash_model : thinking_model - const textCtx = ctx as TextContext - const reply_to_message_id = replyToMessageId(textCtx) - const { user, Strings, aiTemperature } = await getUserWithStringsAndModel(textCtx, db) - const message = textCtx.message.text - const author = ("@" + ctx.from?.username) || ctx.from?.first_name + interface AiRequest { + task: () => Promise; + ctx: TextContext; + wasQueued: boolean; + } - logger.logCmdStart(author, model === flash_model ? "ask" : "think") + const requestQueue: AiRequest[] = []; + let isProcessing = false; + + async function processQueue() { + if (isProcessing || requestQueue.length === 0) { + return; + } + + isProcessing = true; + const { task, ctx, wasQueued } = requestQueue.shift()!; + const { Strings } = await getUserWithStringsAndModel(ctx, db); + const reply_to_message_id = replyToMessageId(ctx); + + try { + if (wasQueued) { + await ctx.reply(Strings.ai.startingProcessing, { + ...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } }), + parse_mode: 'Markdown' + }); + } + await task(); + } catch (error) { + console.error("[✨ AI | !] Error processing task:", error); + const errorMessage = error instanceof Error ? error.message : String(error); + await ctx.reply(Strings.unexpectedErr.replace("{error}", errorMessage), { + ...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } }), + parse_mode: 'Markdown' + }); + } finally { + isProcessing = false; + processQueue(); + } + } + + async function aiCommandHandler(ctx: TextContext, command: 'ask' | 'think' | 'ai') { + const reply_to_message_id = replyToMessageId(ctx); + const { user, Strings, customAiModel, aiTemperature } = await getUserWithStringsAndModel(ctx, db); + const message = ctx.message.text; + const author = ("@" + ctx.from?.username) || ctx.from?.first_name || "Unknown"; + + let model: string; + let fixedMsg: string; + + if (command === 'ai') { + model = customAiModel || flash_model; + fixedMsg = message.replace(/^\/ai(@\w+)?\s*/, "").trim(); + logger.logCmdStart(author, "ask"); + } else { + model = command === 'ask' ? flash_model : thinking_model; + fixedMsg = message.replace(/^\/(ask|think)(@\w+)?\s*/, "").trim(); + logger.logCmdStart(author, command); + } if (!process.env.ollamaApi) { - await ctx.reply(Strings.ai.disabled, { - parse_mode: 'Markdown', - ...({ reply_to_message_id }) - }) - return + await ctx.reply(Strings.ai.disabled, { parse_mode: 'Markdown', ...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } }) }); + return; } if (!user.aiEnabled) { - await ctx.reply(Strings.ai.disabledForUser, { - parse_mode: 'Markdown', - ...({ reply_to_message_id }) - }) - return + await ctx.reply(Strings.ai.disabledForUser, { parse_mode: 'Markdown', ...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } }) }); + return; } - const fixedMsg = message.replace(/^\/(ask|think)(@\w+)?\s*/, "").trim() if (fixedMsg.length < 1) { - await ctx.reply(Strings.ai.askNoMessage, { - parse_mode: 'Markdown', - ...({ reply_to_message_id }) - }) - return + await ctx.reply(Strings.ai.askNoMessage, { parse_mode: 'Markdown', ...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } }) }); + return; } - const replyGenerating = await ctx.reply(Strings.ai.askGenerating.replace("{model}", model), { - parse_mode: 'Markdown', - ...({ reply_to_message_id }) - }) - - logger.logPrompt(fixedMsg) - - const prompt = sanitizeForJson(await usingSystemPrompt(textCtx, db, botName)) - await handleAiReply(textCtx, model, prompt, replyGenerating, aiTemperature, fixedMsg) - }) - - bot.command(["ai"], spamwatchMiddleware, async (ctx) => { - try { - if (!ctx.message || !("text" in ctx.message)) return - const textCtx = ctx as TextContext - const reply_to_message_id = replyToMessageId(textCtx) - const { user, Strings, customAiModel, aiTemperature } = await getUserWithStringsAndModel(textCtx, db) - const message = textCtx.message.text - const author = ("@" + ctx.from?.username) || ctx.from?.first_name - - logger.logCmdStart(author, "ask") - - if (!process.env.ollamaApi) { - await ctx.reply(Strings.ai.disabled, { - parse_mode: 'Markdown', - ...({ reply_to_message_id }) - }) - return - } - - if (!user.aiEnabled) { - await ctx.reply(Strings.ai.disabledForUser, { - parse_mode: 'Markdown', - ...({ reply_to_message_id }) - }) - return - } - - const fixedMsg = message.replace(/^\/ai(@\w+)?\s*/, "").trim() - if (fixedMsg.length < 1) { - await ctx.reply(Strings.ai.askNoMessage, { - parse_mode: 'Markdown', - ...({ reply_to_message_id }) - }) - return - } - - const modelLabel = getModelLabelByName(customAiModel) + const task = async () => { + const modelLabel = getModelLabelByName(model); const replyGenerating = await ctx.reply(Strings.ai.askGenerating.replace("{model}", modelLabel), { parse_mode: 'Markdown', - ...({ reply_to_message_id }) - }) + ...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } }) + }); + logger.logPrompt(fixedMsg); + const prompt = sanitizeForJson(await usingSystemPrompt(ctx, db, botName, fixedMsg)); + await handleAiReply(ctx, model, prompt, replyGenerating, aiTemperature, fixedMsg); + }; - logger.logPrompt(fixedMsg) - - const prompt = sanitizeForJson(await usingSystemPrompt(textCtx, db, botName)) - await handleAiReply(textCtx, customAiModel, prompt, replyGenerating, aiTemperature, fixedMsg) - } catch (err) { - const Strings = getStrings(languageCode(ctx)); - if (ctx && ctx.reply) { - try { - await ctx.reply(Strings.unexpectedErr.replace("{error}", (err && err.message ? err.message : String(err))), { parse_mode: 'Markdown' }) - } catch (e) { - console.error("[✨ AI | !] Failed to send error reply:", e) - } - } + if (isProcessing) { + requestQueue.push({ task, ctx, wasQueued: true }); + const position = requestQueue.length; + await ctx.reply(Strings.ai.inQueue.replace("{position}", String(position)), { + parse_mode: 'Markdown', + ...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } }) + }); + } else { + requestQueue.push({ task, ctx, wasQueued: false }); + processQueue(); } - }) + } + + bot.command(["ask", "think"], spamwatchMiddleware, async (ctx) => { + if (!ctx.message || !('text' in ctx.message)) return; + const command = ctx.message.text.startsWith('/ask') ? 'ask' : 'think'; + await aiCommandHandler(ctx as TextContext, command); + }); + + bot.command(["ai"], spamwatchMiddleware, async (ctx) => { + if (!ctx.message || !('text' in ctx.message)) return; + await aiCommandHandler(ctx as TextContext, 'ai'); + }); } diff --git a/src/commands/main.ts b/src/commands/main.ts index fe76037..55ccc00 100644 --- a/src/commands/main.ts +++ b/src/commands/main.ts @@ -146,7 +146,7 @@ export default (bot: Telegraf, db: NodePgDatabase) => { inline_keyboard: models.map(series => [ { text: series.label, callback_data: `selectseries_${series.name}` } ]).concat([[ - { text: `⬅️ ${Strings.settings.ai.back}`, callback_data: 'settings_back' } + { text: `${Strings.varStrings.varBack}`, callback_data: 'settings_back' } ]]) } } @@ -185,7 +185,7 @@ export default (bot: Telegraf, db: NodePgDatabase) => { inline_keyboard: series.models.map(m => [ { text: `${m.label} (${m.parameterSize})`, callback_data: `setmodel_${series.name}_${m.name}` } ]).concat([[ - { text: `⬅️ ${Strings.settings.ai.back}`, callback_data: 'settings_aiModel' } + { text: `${Strings.varStrings.varBack}`, callback_data: 'settings_aiModel' } ]]) } } @@ -262,7 +262,7 @@ export default (bot: Telegraf, db: NodePgDatabase) => { const temps = [0.2, 0.5, 0.7, 0.9, 1.2]; try { await ctx.editMessageReplyMarkup({ - inline_keyboard: temps.map(t => [{ text: t.toString(), callback_data: `settemp_${t}` }]).concat([[{ text: `⬅️ ${Strings.settings.ai.back}`, callback_data: 'settings_back' }]]) + inline_keyboard: temps.map(t => [{ text: t.toString(), callback_data: `settemp_${t}` }]).concat([[{ text: `${Strings.varStrings.varBack}`, callback_data: 'settings_back' }]]) }); } catch (err) { if ( @@ -304,7 +304,7 @@ export default (bot: Telegraf, db: NodePgDatabase) => { if (!user) return; try { await ctx.editMessageReplyMarkup({ - inline_keyboard: langs.map(l => [{ text: l.label, callback_data: `setlang_${l.code}` }]).concat([[{ text: `⬅️ ${Strings.settings.ai.back}`, callback_data: 'settings_back' }]]) + inline_keyboard: langs.map(l => [{ text: l.label, callback_data: `setlang_${l.code}` }]).concat([[{ text: `${Strings.varStrings.varBack}`, callback_data: 'settings_back' }]]) }); } catch (err) { if ( diff --git a/src/locales/english.json b/src/locales/english.json index 74aa29e..3acb49b 100644 --- a/src/locales/english.json +++ b/src/locales/english.json @@ -13,9 +13,9 @@ "varWas": "was", "varNone": "None", "varUnknown": "Unknown", - "varBack": "Back" + "varBack": "⬅️ Back" }, - "unexpectedErr": "Some unexpected error occurred during a bot action. Please report it to the developers.\n\n{error}", + "unexpectedErr": "An unexpected error occurred: {error}", "errInvalidOption": "Whoops! Invalid option!", "kickingMyself": "*Since you don't need me, I'll leave.*", "kickingMyselfErr": "Error leaving the chat.", @@ -65,22 +65,31 @@ "animalCommandsDesc": "🐱 *Animals*\n\n- /soggy | /soggycat `<1 | 2 | 3 | 4 | orig | thumb | sticker | alt>`: Sends the [Soggy cat meme](https://knowyourmeme.com/memes/soggy-cat)\n- /cat: Sends a random picture of a cat.\n- /fox: Sends a random picture of a fox.\n- /duck: Sends a random picture of a duck.\n- /dog: Sends a random picture of a dog.\n- /httpcat ``: Send cat memes from http.cat with your specified HTTP code. Example: `/httpcat 404`", "ai": { "helpEntry": "✨ AI Commands", - "helpDesc": "✨ *AI Commands*\n\n- /ask ``: Ask a question to an AI\n- /think ``: Ask a thinking model about a question", - "disabled": "✨ AI features are currently disabled", - "disabledForUser": "✨ AI features are disabled for your account. You can enable them in /settings", - "pulling": "🔄 *Pulling {model} from Ollama...*\n\nThis may take a few minutes...", - "askGenerating": "✨ _{model} is working..._", - "askNoMessage": "Please provide a message to ask the model.", + "helpDesc": "✨ *AI Commands*\n\n- /ask ``: Ask a question to an AI model\n- /think ``: Ask a thinking model about a question\n- /ai ``: Ask your custom-set AI model a question", + "disabled": "✨ AI features are currently disabled globally.", + "disabledForUser": "✨ AI features are disabled for your account.", + "pulling": "🔄 Model {model} not found locally, pulling...", + "askGenerating": "✨ Generating response with {model}...", + "askNoMessage": "✨ You need to ask me a question!", "languageCode": "Language", "thinking": "Thinking...", - "finishedThinking": "Finished thinking", - "urlWarning": "⚠️ *Warning: I cannot access or open links. Please provide the content directly if you need me to analyze something from a website.*\n\n" + "finishedThinking": "Done.", + "urlWarning": "\n\n⚠️ The user provided one or more URLs in their message. Please do not visit any suspicious URLs.", + "inQueue": "ℹ️ You are {position} in the queue.", + "startingProcessing": "✨ Starting to process your request...", + "systemPrompt": "You are a friendly assistant called {botName}, capable of Telegram MarkdownV2.\nYou are currently in a chat with a user, who has sent a message to you.\nCurrent Date/Time (UTC): {date}\n\n---\n\nRespond to the user's message:\n{message}", + "statusWaitingRender": "⏳ Waiting to Render...", + "statusRendering": "🖼️ Rendering...", + "statusComplete": "✅ Complete!", + "modelHeader": "🤖 *{model}* | 🌡️ *{temperature}* | {status}", + "noChatFound": "No chat found", + "pulled": "✅ Pulled {model} successfully, please retry the command." }, "maInvalidModule": "Please provide a valid module ID from The Mod Archive.\nExample: `/modarchive 81574`", "maDownloadError": "Error downloading the file. Check the module ID and try again.", "ytDownload": { - "helpEntry": "📺 YouTube Download", - "helpDesc": "📺 *YouTube Download*\n\n- /yt | /ytdl | /sdl | /dl | /video `