From 62247b81bc775ea915eb6d713729529c1de0f190 Mon Sep 17 00:00:00 2001 From: pinks Date: Sun, 24 Sep 2023 21:58:09 +0200 Subject: [PATCH] feat: broadcast command --- bot/broadcastCommand.ts | 71 +++++++++++++++++++++++++++++++++++++++++ bot/cancelCommand.ts | 4 ++- bot/img2imgCommand.ts | 36 ++++++++++++++------- bot/mod.ts | 5 +-- bot/pnginfoCommand.ts | 8 +++-- bot/queueCommand.ts | 5 ++- bot/txt2imgCommand.ts | 28 ++++++++++------ 7 files changed, 130 insertions(+), 27 deletions(-) create mode 100644 bot/broadcastCommand.ts diff --git a/bot/broadcastCommand.ts b/bot/broadcastCommand.ts new file mode 100644 index 0000000..54edcc3 --- /dev/null +++ b/bot/broadcastCommand.ts @@ -0,0 +1,71 @@ +import { CommandContext } from "grammy"; +import { bold, fmt, FormattedString } from "grammy_parse_mode"; +import { distinctBy } from "std/collections"; +import { getConfig } from "../app/config.ts"; +import { generationStore } from "../app/generationStore.ts"; +import { ErisContext, logger } from "./mod.ts"; +import { formatUserChat } from "../utils/formatUserChat.ts"; + +export async function broadcastCommand(ctx: CommandContext) { + if (!ctx.from?.username) { + return ctx.reply("I don't know who you are."); + } + + const config = await getConfig(); + + if (!config.adminUsernames.includes(ctx.from.username)) { + return ctx.reply("Only a bot admin can use this command."); + } + + const text = ctx.match.trim(); + + if (!text) { + return ctx.reply("Please specify a message to broadcast."); + } + + // find users who interacted with bot in the last 24 hours + const gens = await generationStore.getAll( + { after: new Date(Date.now() - 24 * 60 * 60 * 1000) }, + { reverse: true }, + ).then((gens) => distinctBy(gens, (gen) => gen.value.from.id)); + + let sentCount = 0; + const errors: FormattedString[] = []; + const getMessage = () => + fmt([ + fmt`Broadcasted to ${sentCount}/${gens.length} users.\n\n`, + errors.length > 0 ? fmt([bold("Errors:"), "\n", ...errors]) : "", + ]); + + const replyMessage = await ctx.replyFmt(getMessage(), { + reply_to_message_id: ctx.message?.message_id, + }); + + // send message to each user + for (const gen of gens) { + try { + await ctx.api.sendMessage(gen.value.from.id, text); + logger().info(`Broadcasted to ${formatUserChat({ from: gen.value.from })}`); + sentCount++; + } catch (err) { + logger().error(`Broadcasting to ${formatUserChat({ from: gen.value.from })} failed: ${err}`); + errors.push(fmt`${bold(formatUserChat({ from: gen.value.from }))} - ${err.message}\n`); + } + const fmtMessage = getMessage(); + if (sentCount % 20 === 0) { + await ctx.api.editMessageText( + replyMessage.chat.id, + replyMessage.message_id, + fmtMessage.text, + { entities: fmtMessage.entities }, + ).catch(() => undefined); + } + } + const fmtMessage = getMessage(); + await ctx.api.editMessageText( + replyMessage.chat.id, + replyMessage.message_id, + fmtMessage.text, + { entities: fmtMessage.entities }, + ).catch(() => undefined); +} diff --git a/bot/cancelCommand.ts b/bot/cancelCommand.ts index 6a1236b..657992d 100644 --- a/bot/cancelCommand.ts +++ b/bot/cancelCommand.ts @@ -7,5 +7,7 @@ export async function cancelCommand(ctx: ErisContext) { .filter((job) => job.lockUntil < new Date()) .filter((j) => j.state.from.id === ctx.from?.id); for (const job of userJobs) await generationQueue.deleteJob(job.id); - await ctx.reply(`Cancelled ${userJobs.length} jobs`); + await ctx.reply(`Cancelled ${userJobs.length} jobs`, { + reply_to_message_id: ctx.message?.message_id, + }); } diff --git a/bot/img2imgCommand.ts b/bot/img2imgCommand.ts index 12c7cb8..b4f9d4d 100644 --- a/bot/img2imgCommand.ts +++ b/bot/img2imgCommand.ts @@ -28,30 +28,34 @@ async function img2img( state: QuestionState = {}, ): Promise { if (!ctx.message?.from?.id) { - await ctx.reply("I don't know who you are"); + await ctx.reply("I don't know who you are", { + reply_to_message_id: ctx.message?.message_id, + }); return; } const config = await getConfig(); if (config.pausedReason != null) { - await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`); + await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`, { + reply_to_message_id: ctx.message?.message_id, + }); return; } const jobs = await generationQueue.getAllJobs(); if (jobs.length >= config.maxJobs) { - await ctx.reply( - `The queue is full. Try again later. (Max queue size: ${config.maxJobs})`, - ); + await ctx.reply(`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`, { + reply_to_message_id: ctx.message?.message_id, + }); return; } const userJobs = jobs.filter((job) => job.state.from.id === ctx.message?.from?.id); if (userJobs.length >= config.maxUserJobs) { - await ctx.reply( - `You already have ${config.maxUserJobs} jobs in queue. Try again later.`, - ); + await ctx.reply(`You already have ${config.maxUserJobs} jobs in queue. Try again later.`, { + reply_to_message_id: ctx.message?.message_id, + }); return; } @@ -91,7 +95,11 @@ async function img2img( await ctx.reply( "Please show me a picture to repaint." + img2imgQuestion.messageSuffixMarkdown(JSON.stringify(state satisfies QuestionState)), - { reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" }, + { + reply_markup: { force_reply: true, selective: true }, + parse_mode: "Markdown", + reply_to_message_id: ctx.message?.message_id, + }, ); return; } @@ -100,12 +108,18 @@ async function img2img( await ctx.reply( "Please describe the picture you want to repaint." + img2imgQuestion.messageSuffixMarkdown(JSON.stringify(state satisfies QuestionState)), - { reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" }, + { + reply_markup: { force_reply: true, selective: true }, + parse_mode: "Markdown", + reply_to_message_id: ctx.message?.message_id, + }, ); return; } - const replyMessage = await ctx.reply("Accepted. You are now in queue."); + const replyMessage = await ctx.reply("Accepted. You are now in queue.", { + reply_to_message_id: ctx.message?.message_id, + }); await generationQueue.pushJob({ task: { type: "img2img", fileId: state.fileId, params: state.params }, diff --git a/bot/mod.ts b/bot/mod.ts index c4c56f9..653f939 100644 --- a/bot/mod.ts +++ b/bot/mod.ts @@ -1,10 +1,10 @@ import { Api, Bot, Context, RawApi, session, SessionFlavor } from "grammy"; -import { autoQuote } from "grammy_autoquote"; import { FileFlavor, hydrateFiles } from "grammy_files"; import { hydrateReply, ParseModeFlavor } from "grammy_parse_mode"; import { getLogger } from "std/log"; import { getConfig, setConfig } from "../app/config.ts"; import { formatUserChat } from "../utils/formatUserChat.ts"; +import { broadcastCommand } from "./broadcastCommand.ts"; import { cancelCommand } from "./cancelCommand.ts"; import { img2imgCommand, img2imgQuestion } from "./img2imgCommand.ts"; import { pnginfoCommand, pnginfoQuestion } from "./pnginfoCommand.ts"; @@ -45,7 +45,6 @@ export const bot = new Bot( }, ); -bot.use(autoQuote); bot.use(hydrateReply); bot.use(session({ type: "multi", @@ -129,6 +128,8 @@ bot.command("queue", queueCommand); bot.command("cancel", cancelCommand); +bot.command("broadcast", broadcastCommand); + bot.command("pause", async (ctx) => { if (!ctx.from?.username) return; const config = await getConfig(); diff --git a/bot/pnginfoCommand.ts b/bot/pnginfoCommand.ts index 589886a..1294917 100644 --- a/bot/pnginfoCommand.ts +++ b/bot/pnginfoCommand.ts @@ -23,7 +23,11 @@ async function pnginfo(ctx: ErisContext, includeRepliedTo: boolean): Promise) { let formattedMessage = await getMessageText(); - const queueMessage = await ctx.replyFmt(formattedMessage, { disable_notification: true }); + const queueMessage = await ctx.replyFmt(formattedMessage, { + disable_notification: true, + reply_to_message_id: ctx.message?.message_id, + }); handleFutureUpdates().catch(() => undefined); async function getMessageText() { diff --git a/bot/txt2imgCommand.ts b/bot/txt2imgCommand.ts index 535ccad..c4f10a0 100644 --- a/bot/txt2imgCommand.ts +++ b/bot/txt2imgCommand.ts @@ -20,30 +20,32 @@ export async function txt2imgCommand(ctx: CommandContext) { async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolean): Promise { if (!ctx.message?.from?.id) { - await ctx.reply("I don't know who you are"); + await ctx.reply("I don't know who you are", { reply_to_message_id: ctx.message?.message_id }); return; } const config = await getConfig(); if (config.pausedReason != null) { - await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`); + await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`, { + reply_to_message_id: ctx.message?.message_id, + }); return; } const jobs = await generationQueue.getAllJobs(); if (jobs.length >= config.maxJobs) { - await ctx.reply( - `The queue is full. Try again later. (Max queue size: ${config.maxJobs})`, - ); + await ctx.reply(`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`, { + reply_to_message_id: ctx.message?.message_id, + }); return; } const userJobs = jobs.filter((job) => job.state.from.id === ctx.message?.from?.id); if (userJobs.length >= config.maxUserJobs) { - await ctx.reply( - `You already have ${config.maxUserJobs} jobs in queue. Try again later.`, - ); + await ctx.reply(`You already have ${config.maxUserJobs} jobs in queue. Try again later.`, { + reply_to_message_id: ctx.message?.message_id, + }); return; } @@ -69,12 +71,18 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea await ctx.reply( "Please tell me what you want to see." + txt2imgQuestion.messageSuffixMarkdown(), - { reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" }, + { + reply_markup: { force_reply: true, selective: true }, + parse_mode: "Markdown", + reply_to_message_id: ctx.message?.message_id, + }, ); return; } - const replyMessage = await ctx.reply("Accepted. You are now in queue."); + const replyMessage = await ctx.reply("Accepted. You are now in queue.", { + reply_to_message_id: ctx.message?.message_id, + }); await generationQueue.pushJob({ task: { type: "txt2img", params },