Initial commit

src/bot.ts

@@ -0,0 +1,99 @@
+import { logger } from "#root/logger.js";
+import { predict } from "#root/predict.js";
+import { downloadImage, isFooocosAlive, uploadPhoto } from "#root/utils.js";
+import { apiThrottler } from "@grammyjs/transformer-throttler";
+import { Bot, Context, InlineKeyboard, InlineQueryResultBuilder, InputMediaBuilder } from "grammy";
+
+const activeTasks = new Map<number, boolean>();
+
+const articlePhoto = {
+    url: "https://i.imgur.com/oAWNcVt.jpg",
+    width: 512,
+    height: 512
+};
+
+async function handleInlineQuery(ctx: Context) {
+    if (!await isFooocosAlive()) {
+        const answer = InlineQueryResultBuilder.article('fooocos:down', "Сервис недоступен")
+            .text('Сервис временно недоступен. Попробуйте позже');
+        return ctx.answerInlineQuery([answer], { cache_time: 0 });
+    }
+
+    const userId = Number(ctx.update.inline_query?.from.id);
+    const prompt = ctx.inlineQuery?.query;
+
+    if (activeTasks.get(userId)) {
+        logger.trace(`${userId} in activeTasks`);
+        return ctx.answerInlineQuery([]);
+    }
+
+    const answer = prompt ?
+        InlineQueryResultBuilder.photo(`fooocos:query:prompt`, articlePhoto.url, {
+            caption: prompt,
+            thumbnail_url: articlePhoto.url,
+            photo_width: articlePhoto.width,
+            photo_height: articlePhoto.height,
+            reply_markup: new InlineKeyboard()
+                .switchInlineCurrent('Вы в очереди')
+        }) :
+        InlineQueryResultBuilder.article('fooocos:query', 'введите промпт', {
+            reply_markup: new InlineKeyboard()
+                .switchInlineCurrent('Попробовать снова')
+        })
+            .text('Нет промпта');
+
+    return ctx.answerInlineQuery([answer], { cache_time: 0 });
+}
+
+async function handleChosenInlineResult(ctx: Context) {
+    const userId = ctx.update.chosen_inline_result?.from.id!;
+    const messageId = ctx.chosenInlineResult?.inline_message_id!;
+    const prompt = ctx.chosenInlineResult?.query!;
+
+    if (!prompt) {
+        return ctx.api.editMessageTextInline(messageId, "Ошибка: нет промпта. Попробуй ещё раз!");
+    }
+
+    activeTasks.set(userId, true);
+
+    const gen = predict(prompt);
+
+    for await (const state of gen) {
+        if (state.status === "RUNNING") {
+            if (state.preview) {
+                const previewImage = Buffer.from(state.preview, 'base64');
+                const fileId = await uploadPhoto(ctx, previewImage);
+                const media = InputMediaBuilder.photo(fileId, {
+                    caption: `Промпт: ${prompt}`
+                });
+
+                try {
+                    await ctx.api.editMessageMediaInline(messageId, media);
+                } catch (err) {
+                    activeTasks.delete(userId);
+                    return;
+                }
+            }
+        } else if (state.status === "SUCCESS") {
+            const resultImage = await downloadImage(state.result![0]);
+            const fileId = await uploadPhoto(ctx, resultImage);
+            const media = InputMediaBuilder.photo(fileId, {
+                caption: `Промпт: ${prompt}`
+            });
+            activeTasks.delete(userId);
+            return ctx.api.editMessageMediaInline(messageId, media);
+        } else if (state.status === 'ERROR') {
+            activeTasks.delete(userId);
+            logger.error('ERROR: failed to generate image');
+        }
+    }
+}
+
+export function createBot(token: string) {
+    const bot = new Bot(token);
+    bot.api.config.use(apiThrottler());
+    const protectedBot = bot.errorBoundary((err) => logger.error(err));
+    protectedBot.on('chosen_inline_result', handleChosenInlineResult);
+    protectedBot.on('inline_query', handleInlineQuery);
+    return bot;
+}