forked from pinks/eris
1
0
Fork 0
eris/bot/mod.ts

160 lines
4.9 KiB
TypeScript
Raw Normal View History

2023-09-24 13:08:35 +00:00
import { Api, Bot, Context, RawApi, session, SessionFlavor } from "grammy";
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";
2023-09-23 18:49:05 +00:00
import { formatUserChat } from "../utils/formatUserChat.ts";
2023-09-24 19:58:09 +00:00
import { broadcastCommand } from "./broadcastCommand.ts";
2023-09-24 13:08:35 +00:00
import { cancelCommand } from "./cancelCommand.ts";
import { img2imgCommand, img2imgQuestion } from "./img2imgCommand.ts";
import { pnginfoCommand, pnginfoQuestion } from "./pnginfoCommand.ts";
2023-09-10 18:56:17 +00:00
import { queueCommand } from "./queueCommand.ts";
2023-09-10 23:59:33 +00:00
import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts";
2023-09-10 18:56:17 +00:00
2023-09-24 13:08:35 +00:00
export const logger = () => getLogger();
2023-09-10 18:56:17 +00:00
2023-09-22 02:59:22 +00:00
interface SessionData {
2023-09-24 13:08:35 +00:00
chat: ErisChatData;
user: ErisUserData;
2023-09-22 02:59:22 +00:00
}
2023-09-24 13:08:35 +00:00
interface ErisChatData {
2023-09-22 02:59:22 +00:00
language?: string;
}
2023-09-24 13:08:35 +00:00
interface ErisUserData {
2023-09-22 02:59:22 +00:00
params?: Record<string, string>;
}
2023-09-24 13:08:35 +00:00
export type ErisContext =
& FileFlavor<ParseModeFlavor<Context>>
& SessionFlavor<SessionData>;
2023-09-22 02:59:22 +00:00
2023-09-24 13:08:35 +00:00
type WithRetryApi<T extends RawApi> = {
2023-09-11 17:07:46 +00:00
[M in keyof T]: T[M] extends (args: infer P, ...rest: infer A) => infer R
? (args: P extends object ? P & { maxAttempts?: number; maxWait?: number } : P, ...rest: A) => R
2023-09-11 17:07:46 +00:00
: T[M];
};
2023-09-24 13:08:35 +00:00
type ErisApi = Api<WithRetryApi<RawApi>>;
2023-09-22 02:59:22 +00:00
2023-09-24 13:08:35 +00:00
export const bot = new Bot<ErisContext, ErisApi>(
Deno.env.get("TG_BOT_TOKEN")!,
{
client: { timeoutSeconds: 20 },
},
);
2023-09-22 02:59:22 +00:00
2023-09-24 13:08:35 +00:00
bot.use(hydrateReply);
bot.use(session<SessionData, ErisContext>({
2023-09-22 02:59:22 +00:00
type: "multi",
chat: {
initial: () => ({}),
},
user: {
getSessionKey: (ctx) => ctx.from?.id.toFixed(),
initial: () => ({}),
},
}));
2023-09-10 18:56:17 +00:00
2023-09-24 13:08:35 +00:00
bot.api.config.use(hydrateFiles(bot.token));
2023-09-10 18:56:17 +00:00
// Automatically retry bot requests if we get a "too many requests" or telegram internal error
bot.api.config.use(async (prev, method, payload, signal) => {
2023-09-11 17:07:46 +00:00
const maxAttempts = payload && ("maxAttempts" in payload) ? payload.maxAttempts ?? 3 : 3;
const maxWait = payload && ("maxWait" in payload) ? payload.maxWait ?? 10 : 10;
2023-09-10 18:56:17 +00:00
let attempt = 0;
while (true) {
attempt++;
const result = await prev(method, payload, signal);
if (result.ok) return result;
if (result.error_code !== 429) return result;
if (attempt >= maxAttempts) return result;
const retryAfter = result.parameters?.retry_after ?? (attempt * 5);
if (retryAfter > maxWait) return result;
2023-09-11 17:07:46 +00:00
logger().warning(
2023-09-11 20:48:38 +00:00
`${method} (attempt ${attempt}) failed: ${result.error_code} ${result.description}`,
2023-09-11 17:07:46 +00:00
);
await new Promise((resolve) => setTimeout(resolve, retryAfter * 1000));
2023-09-10 18:56:17 +00:00
}
});
2023-09-11 20:43:12 +00:00
bot.catch((err) => {
2023-09-11 20:48:38 +00:00
logger().error(
`Handling update from ${formatUserChat(err.ctx)} failed: ${err.name} ${err.message}`,
);
2023-09-11 20:43:12 +00:00
});
2023-09-10 18:56:17 +00:00
// if error happened, try to reply to the user with the error
bot.use(async (ctx, next) => {
try {
await next();
} catch (err) {
try {
await ctx.reply(`Handling update failed: ${err}`, {
reply_to_message_id: ctx.message?.message_id,
});
} catch {
throw err;
}
}
});
bot.api.setMyShortDescription("I can generate furry images from text");
bot.api.setMyDescription(
"I can generate furry images from text. " +
"Send /txt2img to generate an image.",
);
bot.api.setMyCommands([
2023-09-18 15:51:19 +00:00
{ command: "txt2img", description: "Generate image from text" },
{ command: "img2img", description: "Generate image from image" },
2023-09-11 20:43:12 +00:00
{ command: "pnginfo", description: "Show generation parameters of an image" },
2023-09-10 18:56:17 +00:00
{ command: "queue", description: "Show the current queue" },
2023-09-18 15:51:19 +00:00
{ command: "cancel", description: "Cancel all your requests" },
2023-09-10 18:56:17 +00:00
]);
bot.command("start", (ctx) => ctx.reply("Hello! Use the /txt2img command to generate an image"));
bot.command("txt2img", txt2imgCommand);
2023-09-11 16:03:07 +00:00
bot.use(txt2imgQuestion.middleware());
2023-09-12 01:57:44 +00:00
bot.command("img2img", img2imgCommand);
bot.use(img2imgQuestion.middleware());
2023-09-10 18:56:17 +00:00
2023-09-11 20:43:12 +00:00
bot.command("pnginfo", pnginfoCommand);
bot.use(pnginfoQuestion.middleware());
2023-09-10 18:56:17 +00:00
bot.command("queue", queueCommand);
2023-09-18 15:51:19 +00:00
bot.command("cancel", cancelCommand);
2023-09-24 19:58:09 +00:00
bot.command("broadcast", broadcastCommand);
2023-09-22 02:59:22 +00:00
bot.command("pause", async (ctx) => {
2023-09-10 18:56:17 +00:00
if (!ctx.from?.username) return;
2023-09-22 02:59:22 +00:00
const config = await getConfig();
2023-09-10 18:56:17 +00:00
if (!config.adminUsernames.includes(ctx.from.username)) return;
if (config.pausedReason != null) {
return ctx.reply(`Already paused: ${config.pausedReason}`);
}
config.pausedReason = ctx.match ?? "No reason given";
2023-09-22 02:59:22 +00:00
await setConfig(config);
2023-09-10 18:56:17 +00:00
logger().warning(`Bot paused by ${ctx.from.first_name} because ${config.pausedReason}`);
return ctx.reply("Paused");
});
2023-09-22 02:59:22 +00:00
bot.command("resume", async (ctx) => {
2023-09-10 18:56:17 +00:00
if (!ctx.from?.username) return;
2023-09-22 02:59:22 +00:00
const config = await getConfig();
2023-09-10 18:56:17 +00:00
if (!config.adminUsernames.includes(ctx.from.username)) return;
if (config.pausedReason == null) return ctx.reply("Already running");
config.pausedReason = null;
2023-09-22 02:59:22 +00:00
await setConfig(config);
2023-09-10 18:56:17 +00:00
logger().info(`Bot resumed by ${ctx.from.first_name}`);
return ctx.reply("Resumed");
});
bot.command("crash", () => {
throw new Error("Crash command used");
});