Compare commits

...

3 Commits

Author SHA1 Message Date
pinks 4d69b87b97 small changes 2023-09-11 19:08:12 +02:00
pinks 6b0d7b7198 better tg api retry logic 2023-09-11 19:07:46 +02:00
pinks 58e8e5c6ea deps 2023-09-11 18:03:07 +02:00
7 changed files with 34 additions and 19 deletions

View File

@ -6,8 +6,16 @@ import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts";
export const logger = () => Log.getLogger();
type WithRetryApi<T extends Grammy.RawApi> = {
[M in keyof T]: T[M] extends (args: infer P, ...rest: infer A) => infer R
? (args: P extends object ? P & { maxAttempts?: number } : P, ...rest: A) => R
: T[M];
};
export type Context = GrammyParseMode.ParseModeFlavor<Grammy.Context> & SessionFlavor;
export const bot = new Grammy.Bot<Context>(Deno.env.get("TG_BOT_TOKEN") ?? "");
export const bot = new Grammy.Bot<Context, Grammy.Api<WithRetryApi<Grammy.RawApi>>>(
Deno.env.get("TG_BOT_TOKEN") ?? "",
);
bot.use(GrammyAutoQuote.autoQuote);
bot.use(GrammyParseMode.hydrateReply);
bot.use(session);
@ -18,6 +26,7 @@ bot.catch((err) => {
// 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;
let attempt = 0;
while (true) {
attempt++;
@ -25,10 +34,13 @@ bot.api.config.use(async (prev, method, payload, signal) => {
if (
result.ok ||
![429, 500].includes(result.error_code) ||
attempt >= 5
attempt >= maxAttempts
) {
return result;
}
logger().warning(
`Retrying ${method} after attempt ${attempt} failed with ${result.error_code} error`,
);
const retryAfterMs = (result.parameters?.retry_after ?? (attempt * 5)) * 1000;
await new Promise((resolve) => setTimeout(resolve, retryAfterMs));
}
@ -62,7 +74,7 @@ bot.api.setMyCommands([
bot.command("start", (ctx) => ctx.reply("Hello! Use the /txt2img command to generate an image"));
bot.command("txt2img", txt2imgCommand);
bot.use(txt2imgQuestion.middleware() as any);
bot.use(txt2imgQuestion.middleware());
bot.command("queue", queueCommand);

View File

@ -15,7 +15,6 @@ export async function queueCommand(ctx: Grammy.CommandContext<Context>) {
const waitingJobs = await jobStore.getBy("status.type", "waiting")
.then((jobs) => jobs.map((job, index) => ({ ...job.value, place: index + 1 })));
const jobs = [...processingJobs, ...waitingJobs];
const config = ctx.session.global;
const { bold } = GrammyParseMode;
return fmt([
"Current queue:\n",
@ -40,7 +39,7 @@ export async function queueCommand(ctx: Grammy.CommandContext<Context>) {
])
: ["Queue is empty.\n"],
"\nActive workers:\n",
...config.workers.flatMap((worker) => [
...ctx.session.global.workers.flatMap((worker) => [
runningWorkers.has(worker.name) ? "✅ " : "☠️ ",
fmt`${bold(worker.name)} `,
`(max ${(worker.maxResolution / 1000000).toFixed(1)} Mpx) `,

View File

@ -4,11 +4,11 @@ import { jobStore } from "../db/jobStore.ts";
import { parsePngInfo } from "../sd.ts";
import { Context, logger } from "./mod.ts";
export const txt2imgQuestion = new GrammyStatelessQ.StatelessQuestion(
export const txt2imgQuestion = new GrammyStatelessQ.StatelessQuestion<Context>(
"txt2img",
async (ctx) => {
if (!ctx.message.text) return;
await txt2img(ctx as any, ctx.message.text, false);
await txt2img(ctx, ctx.message.text, false);
},
);
@ -22,24 +22,23 @@ async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean):
return;
}
const config = ctx.session.global;
if (config.pausedReason != null) {
await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`);
if (ctx.session.global.pausedReason != null) {
await ctx.reply(`I'm paused: ${ctx.session.global.pausedReason || "No reason given"}`);
return;
}
const jobs = await jobStore.getBy("status.type", "waiting");
if (jobs.length >= config.maxJobs) {
if (jobs.length >= ctx.session.global.maxJobs) {
await ctx.reply(
`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`,
`The queue is full. Try again later. (Max queue size: ${ctx.session.global.maxJobs})`,
);
return;
}
const userJobs = jobs.filter((job) => job.value.request.from.id === ctx.message?.from?.id);
if (userJobs.length >= config.maxUserJobs) {
if (userJobs.length >= ctx.session.global.maxUserJobs) {
await ctx.reply(
`You already have ${config.maxUserJobs} jobs in queue. Try again later.`,
`You already have ${ctx.session.global.maxUserJobs} jobs in queue. Try again later.`,
);
return;
}
@ -48,6 +47,7 @@ async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean):
const repliedToMsg = ctx.message.reply_to_message;
const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
if (includeRepliedTo && repliedToText) {
// TODO: remove bot command from replied to text
const originalParams = parsePngInfo(repliedToText);
params = {
...originalParams,

View File

@ -11,7 +11,7 @@ export * as GrammyTypes from "https://deno.land/x/grammy_types@v3.2.0/mod.ts";
export * as GrammyAutoQuote from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts";
export * as GrammyParseMode from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
export * as GrammyKvStorage from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts";
export * as GrammyStatelessQ from "npm:@grammyjs/stateless-question";
export * as GrammyStatelessQ from "https://deno.land/x/grammy_stateless_question_alpha@v3.0.3/mod.ts";
export * as FileType from "npm:file-type@18.5.0";
// @deno-types="./types/png-chunks-extract.d.ts"
export * as PngChunksExtract from "npm:png-chunks-extract@1.0.0";

View File

@ -121,6 +121,7 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
`Generating your prompt now... ${
(progress.progress * 100).toFixed(0)
}% using ${worker.name}`,
{ maxAttempts: 1 },
).catch(() => undefined);
}
await job.update({
@ -130,7 +131,7 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
worker: worker.name,
updatedDate: new Date(),
},
}).catch(() => undefined);
}, { maxAttempts: 1 }).catch(() => undefined);
},
);
@ -140,6 +141,7 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
job.value.reply.chat.id,
job.value.reply.message_id,
`Uploading your images...`,
{ maxAttempts: 1 },
).catch(() => undefined);
}
@ -185,6 +187,7 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
// send the result to telegram
const resultMessage = await bot.api.sendMediaGroup(job.value.request.chat.id, inputFiles, {
reply_to_message_id: job.value.request.message_id,
maxAttempts: 5,
});
// send caption in separate message if it couldn't fit
if (caption.text.length > 1024 && caption.text.length <= 4096) {

View File

@ -19,6 +19,7 @@ export async function updateJobStatusMsgs(): Promise<never> {
job.value.reply.chat.id,
job.value.reply.message_id,
`You are ${formatOrdinal(index + 1)} in queue.`,
{ maxAttempts: 1 },
).catch(() => undefined);
}
} catch (err) {

View File

@ -105,7 +105,7 @@ const languageToFlagMap: Record<string, string> = {
"lb": "🇱🇺", // Luxembourgish - Luxembourg
};
export function getFlagEmoji(countryCode?: string): string | undefined {
if (!countryCode) return;
return languageToFlagMap[countryCode];
export function getFlagEmoji(languageCode?: string): string | undefined {
if (!languageCode) return;
return languageToFlagMap[languageCode];
}