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

229 lines
6.8 KiB
TypeScript

import { Api, Bot, Context, RawApi, session, SessionFlavor } from "grammy";
import { FileFlavor, hydrateFiles } from "grammy_files";
import { hydrateReply, ParseModeFlavor } from "grammy_parse_mode";
import { 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";
import { setConfig, getConfig } from "../app/config.ts";
// Set the new configuration
await setConfig({ maxUserJobs: 1, maxJobs: 500 });
// Fetch the updated configuration
const updatedConfig = await getConfig();
// Log the updated configuration to the console
console.log("Updated Configuration:", updatedConfig);
interface SessionData {
chat: ErisChatData;
user: ErisUserData;
}
interface ErisChatData {
language?: string | undefined;
}
interface ErisUserData {
params?: Record<string, string> | undefined;
}
export type ErisContext =
& FileFlavor<ParseModeFlavor<Context>>
& SessionFlavor<SessionData>;
type WithRetryApi<T extends RawApi> = {
[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<WithRetryApi<RawApi>>;
export const bot = new Bot<ErisContext, ErisApi>(
Deno.env.get("TG_BOT_TOKEN")!,
{
client: { timeoutSeconds: 20 },
},
);
bot.use(hydrateReply);
bot.use(sequentialize((ctx) => ctx.chat?.id.toString()));
bot.use(session<SessionData, ErisContext>({
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;
}
}
});
// Wrap the calls in try-catch for error handling
async function setupBotCommands() {
try {
await bot.api.setMyShortDescription("Generate furry images in '704x704', '576x832', '832x576'. https://ko-fi.com/nyxthebot https://nyx.akiru.de/");
} catch (err) {
error(`Failed to set short description: ${err.message}`);
}
try {
await bot.api.setMyDescription(
"I can generate furry images from text. If you want a different size, use 'Size: 576x832' for example." +
"Send /txt2img to generate an image.",
);
} catch (err) {
error(`Failed to set description: ${err.message}`);
}
try {
await 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" },
]);
} catch (err) {
error(`Failed to set commands: ${err.message}`);
}
}
// Call the setup function
setupBotCommands();
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");
});
// Set up the webhook in the telegram API and initialize the bot
await bot.api.setWebhook('https://nyx.akiru.de/webhook');
await bot.init();
// Function to handle incoming webhook requests
export async function handleWebhook(req: Request): Promise<Response> {
try {
const body = await req.json();
// console.log("Received webhook data:", JSON.stringify(body, null, 2));
// Log before processing update
// console.log("Processing update through handleUpdate...");
await bot.handleUpdate(body); // Process the update
// Log after processing update
// console.log("Update processed successfully.");
return new Response(JSON.stringify({ status: 'ok' }), { headers: { 'Content-Type': 'application/json' } });
} catch (error) {
// Detailed error logging
console.error("Error in handleWebhook:", error);
return new Response("Error", { status: 500 });
}
}