pinks
/
eris
1
0
Fork 4
Code Issues 15 Pull Requests 1 Activity

feat: admin priority 

ignore queue limits and assign higher priority on admin jobs
This commit is contained in:
nameless 2023-11-10 01:49:06 +00:00
parent 73810d3204
commit ac63e39373
1 changed files with 22 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
bot/txt2imgCommand.ts
View File

@ -6,6 +6,7 @@ import { generationQueue } from "../app/generationQueue.ts";
import { formatUserChat } from "../utils/formatUserChat.ts"; import { formatUserChat } from "../utils/formatUserChat.ts";
import { ErisContext } from "./mod.ts"; import { ErisContext } from "./mod.ts";
import { getPngInfo, parsePngInfo, PngInfo } from "./parsePngInfo.ts"; import { getPngInfo, parsePngInfo, PngInfo } from "./parsePngInfo.ts";
import { adminStore } from "../app/adminStore.ts";
export const txt2imgQuestion = new StatelessQuestion<ErisContext>( export const txt2imgQuestion = new StatelessQuestion<ErisContext>(
"txt2img", "txt2img",
@ -29,6 +30,7 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea
} }
const config = await getConfig(); const config = await getConfig();
let priority = 0;
if (config.pausedReason != null) { 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"}`, {
@ -37,20 +39,26 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea
return; return;
} }
const jobs = await generationQueue.getAllJobs(); const [admin] = await adminStore.getBy("tgUserId", { value: ctx.from.id });
if (jobs.length >= 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 (admin) {
if (userJobs.length >= config.maxUserJobs) { priority = 1;
await ctx.reply(`You already have ${userJobs.length} jobs in queue. Try again later.`, { } else {
reply_to_message_id: ctx.message?.message_id, const jobs = await generationQueue.getAllJobs();
}); if (jobs.length >= config.maxJobs) {
return; 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 ${userJobs.length} jobs in the queue. Try again later.`, {
reply_to_message_id: ctx.message?.message_id,
});
return;
}
} }
let params: Partial<PngInfo> = {}; let params: Partial<PngInfo> = {};
@ -103,7 +111,7 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea
chat: ctx.message.chat, chat: ctx.message.chat,
requestMessage: ctx.message, requestMessage: ctx.message,
replyMessage: replyMessage, replyMessage: replyMessage,
}, { retryCount: 3, retryDelayMs: 10_000 }); }, { retryCount: 3, retryDelayMs: 10_000, priority: priority });
debug(`Generation (txt2img) enqueued for ${formatUserChat(ctx.message)}`); debug(`Generation (txt2img) enqueued for ${formatUserChat(ctx.message)}`);
} }