ABOCN/TelegramBot
1
0
Fork 0
Code Issues 2 Pull requests Projects Releases Packages Wiki Activity

max queue size

Browse source
This commit is contained in:
Aidan 2025-07-03 14:35:59 -04:00
parent 6b1c7f0da9
commit 19e794e34c
4 changed files with 19 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

18
src/commands/ai.ts
View file

@ -41,7 +41,7 @@ import { ensureUserInDb } from "../utils/ensure-user"
import * as schema from '../db/schema'
import type { NodePgDatabase } from "drizzle-orm/node-postgres"
import { eq, sql } from 'drizzle-orm'
import { models, unloadModelAfterB } from "../../config/ai"
import { models, unloadModelAfterB, maxUserQueueSize } from "../../config/ai"
const spamwatchMiddleware = spamwatchMiddlewareModule(isOnSpamWatch)
export const flash_model = process.env.flashModel || "gemma3:4b"
@ -503,6 +503,7 @@ export default (bot: Telegraf<Context>, db: NodePgDatabase<typeof schema>) => {
task: () => Promise<void>;
ctx: TextContext;
wasQueued: boolean;
userId: number;
}
const requestQueue: AiRequest[] = [];
@ -567,6 +568,17 @@ export default (bot: Telegraf<Context>, db: NodePgDatabase<typeof schema>) => {
return;
}
const userId = ctx.from!.id;
const userQueueSize = requestQueue.filter(req => req.userId === userId).length;
if (userQueueSize >= maxUserQueueSize) {
await ctx.reply(Strings.ai.queueFull, {
parse_mode: 'Markdown',
...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } })
});
return;
}
const task = async () => {
const modelLabel = getModelLabelByName(model);
const replyGenerating = await ctx.reply(Strings.ai.askGenerating.replace("{model}", `\`${modelLabel}\``), {
@ -578,14 +590,14 @@ export default (bot: Telegraf<Context>, db: NodePgDatabase<typeof schema>) => {
};
if (isProcessing) {
requestQueue.push({ task, ctx, wasQueued: true });
requestQueue.push({ task, ctx, wasQueued: true, userId: ctx.from!.id });
const position = requestQueue.length;
await ctx.reply(Strings.ai.inQueue.replace("{position}", String(position)), {
parse_mode: 'Markdown',
...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } })
});
} else {
requestQueue.push({ task, ctx, wasQueued: false });
requestQueue.push({ task, ctx, wasQueued: false, userId: ctx.from!.id });
processQueue();
}
}