import { Api, Bot, Context, RawApi, session, SessionFlavor } from "grammy"; import { FileFlavor, hydrateFiles } from "grammy_files"; import { hydrateReply, ParseModeFlavor } from "grammy_parse_mode"; import { run, sequentialize } from "grammy_runner"; import { error, info, warning } from "std/log/mod.ts"; import { sessions } from "../api/sessionsRoute.ts"; import { formatUserChat } from "../utils/formatUserChat.ts"; import { omitUndef } from "../utils/omitUndef.ts"; import { broadcastCommand } from "./broadcastCommand.ts"; import { cancelCommand } from "./cancelCommand.ts"; import { img2imgCommand, img2imgQuestion } from "./img2imgCommand.ts"; import { pnginfoCommand, pnginfoQuestion } from "./pnginfoCommand.ts"; import { queueCommand } from "./queueCommand.ts"; import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts"; interface SessionData { chat: ErisChatData; user: ErisUserData; } interface ErisChatData { language?: string | undefined; } interface ErisUserData { params?: Record | undefined; } export type ErisContext = & FileFlavor> & SessionFlavor; type WithRetryApi = { [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 : T[M]; }; type ErisApi = Api>; export const bot = new Bot( Deno.env.get("TG_BOT_TOKEN")!, { client: { timeoutSeconds: 20 }, }, ); bot.use(hydrateReply); bot.use(sequentialize((ctx) => ctx.chat?.id.toString())); bot.use(session({ type: "multi", chat: { initial: () => ({}), }, user: { getSessionKey: (ctx) => ctx.from?.id.toFixed(), initial: () => ({}), }, })); bot.api.config.use(hydrateFiles(bot.token)); // Automatically retry bot requests if we get a "too many requests" or telegram internal error bot.api.config.use(async (prev, method, payload, signal) => { const maxAttempts = payload && ("maxAttempts" in payload) ? payload.maxAttempts ?? 3 : 3; const maxWait = payload && ("maxWait" in payload) ? payload.maxWait ?? 10 : 10; 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; warning( `${method} (attempt ${attempt}) failed: ${result.error_code} ${result.description}`, ); await new Promise((resolve) => setTimeout(resolve, retryAfter * 1000)); } }); bot.catch((err) => { error( `Handling update from ${formatUserChat(err.ctx)} failed: ${err.name} ${err.message}`, ); }); // 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}`, omitUndef({ reply_to_message_id: ctx.message?.message_id, allow_sending_without_reply: true, }), ); } 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([ { command: "txt2img", description: "Generate image from text" }, { command: "img2img", description: "Generate image from image" }, { command: "pnginfo", description: "Try to extract prompt from raw file" }, { command: "queue", description: "Show the current queue" }, { command: "cancel", description: "Cancel all your requests" }, ]); bot.command("start", async (ctx) => { if (ctx.match) { const id = ctx.match.trim(); const session = sessions.get(id); if (session == null) { await ctx.reply( "Login failed: Invalid session ID", omitUndef({ reply_to_message_id: ctx.message?.message_id, }), ); return; } session.userId = ctx.from?.id; sessions.set(id, session); info(`User ${formatUserChat(ctx)} logged in`); // TODO: show link to web ui await ctx.reply( "Login successful! You can now return to the WebUI.", omitUndef({ reply_to_message_id: ctx.message?.message_id, }), ); return; } await ctx.reply( "Hello! Use the /txt2img command to generate an image", omitUndef({ reply_to_message_id: ctx.message?.message_id, }), ); }); bot.command("txt2img", txt2imgCommand); bot.use(txt2imgQuestion.middleware()); bot.command("img2img", img2imgCommand); bot.use(img2imgQuestion.middleware()); bot.command("pnginfo", pnginfoCommand); bot.use(pnginfoQuestion.middleware()); bot.command("queue", queueCommand); bot.command("cancel", cancelCommand); bot.command("broadcast", broadcastCommand); bot.command("crash", () => { throw new Error("Crash command used"); }); export async function runBot() { const runner = run(bot, { runner: { silent: true } }); await runner.task(); }