Add paginated search

OperKH · Feb 27, 2024 · 6fa9e2a · 6fa9e2a
1 parent 8361f4a
commit 6fa9e2a
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 62 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "ai_bot",
-  "version": "1.2.6",
+  "version": "1.2.7",
   "private": true,
   "main": "src/app.ts",
   "type": "module",

diff --git a/src/bot/commands/mediaTracker.command.ts b/src/bot/commands/mediaTracker.command.ts
@@ -1,6 +1,5 @@
 import { Context, NarrowedContext } from 'telegraf';
-import { Message, Update } from 'telegraf/types';
-import { CommandContextExtn } from 'node_modules/telegraf/typings/telegram-types.js';
+import { CallbackQuery, InlineKeyboardMarkup, Message, Update } from 'telegraf/types';
 import { message } from 'telegraf/filters';
 import { TelegramClient, Api } from 'telegram';
 import { StringSession } from 'telegram/sessions/index.js';
@@ -33,7 +32,29 @@ export class MediaTrackerCommand extends Command {
       }
     });
     this.bot.command(this.command, async (ctx) => {
-      await this.searchHandler(ctx);
+      if (ctx.payload) {
+        this.searchAndReplyPaginated(ctx, ctx.chat.id, ctx.message.message_id, ctx.payload, 0);
+      } else {
+        await ctx.reply(`ℹ️ Додай пошуковий запит після команди, наприклад: /${this.command} ігрова консоль`, {
+          reply_parameters: { message_id: ctx.message.message_id },
+        });
+      }
+    });
+    this.bot.action(/^islm-(.+)$/, async (ctx) => {
+      const payload = JSON.parse(ctx.match[1]) as unknown;
+      if (
+        typeof payload === 'object' &&
+        payload !== null &&
+        't' in payload &&
+        'o' in payload &&
+        typeof payload.t === 'string' &&
+        typeof payload.o === 'number'
+      ) {
+        const { t, o } = payload;
+        const chat = await ctx.getChat();
+        await this.searchAndReplyPaginated(ctx, chat.id, undefined, t, o);
+        await ctx.editMessageReplyMarkup({ inline_keyboard: [] });
+      }
     });
     this.bot.command('starthistoryimport', async (ctx) => {
       this.startHistoryImport(ctx);
@@ -81,8 +102,8 @@ export class MediaTrackerCommand extends Command {
             reply_parameters: { message_id: Number(messageId), allow_sending_without_reply: true },
             disable_notification: true,
           });
-          // Wait 300ms before send next message
-          await new Promise((r) => setTimeout(r, 300));
+          // Wait 1 second before send next message
+          await new Promise((r) => setTimeout(r, 1000));
         }
       }
       // Save to DB
@@ -96,67 +117,80 @@ export class MediaTrackerCommand extends Command {
     }
   }
 
-  private async searchHandler(
-    ctx: Context<{
-      message: Update.New & Update.NonChannel & Message.TextMessage;
-      update_id: number;
-    }> &
-      Omit<IBotContext, keyof Context<Update>> &
-      CommandContextExtn,
+  private async searchAndReplyPaginated(
+    ctx:
+      | NarrowedContext<IBotContext, Update.MessageUpdate<Message>>
+      | Context<Update.CallbackQueryUpdate<CallbackQuery>>,
+    chatId: number,
+    firstMessageId: number | undefined,
+    text: string,
+    offset: number,
   ) {
-    if (ctx.payload) {
-      const chatId = ctx.chat.id;
-      const messageId = ctx.message.message_id;
-      const textEmbedding = await this.aiService.getTextClipEmbedding(ctx.payload);
-      const textEmbeddingString = JSON.stringify(textEmbedding);
-      type Messages = {
-        messageId: string;
-        similarity: number;
-      };
-      const chatPhotoMessageRepository = this.dataSource.getRepository(ChatPhotoMessage);
-      const t1 = performance.now();
-      const messages = await chatPhotoMessageRepository
-        .createQueryBuilder('msg')
-        .select('msg.messageId', 'messageId')
-        .addSelect('1 - (embedding <=> :embedding)', 'similarity')
-        .where('msg.chatId = :chatId')
-        .andWhere('1 - (embedding <=> :embedding) > :matchImageThreshold')
-        .orderBy('similarity', 'DESC')
-        .limit(this.configService.get('MATCH_IMAGE_COUNT'))
-        .setParameters({
-          chatId,
-          embedding: textEmbeddingString,
-          matchImageThreshold: this.configService.get('MATCH_TEXT_THRESHOLD'),
-        })
-        .getRawMany<Messages>();
-      const t2 = performance.now();
-      console.log(`DB query time: ${Math.round(t2 - t1)} ms`);
-      // When similar
-      if (messages.length > 0) {
+    const textEmbedding = await this.aiService.getTextClipEmbedding(text);
+    const textEmbeddingString = JSON.stringify(textEmbedding);
+    type Messages = {
+      messageId: string;
+      similarity: number;
+    };
+    const chatPhotoMessageRepository = this.dataSource.getRepository(ChatPhotoMessage);
+    const t1 = performance.now();
+    const limit = this.configService.get('MATCH_IMAGE_COUNT');
+    const messages = await chatPhotoMessageRepository
+      .createQueryBuilder('msg')
+      .select('msg.messageId', 'messageId')
+      .addSelect('1 - (embedding <=> :embedding)', 'similarity')
+      .where('msg.chatId = :chatId')
+      .andWhere('1 - (embedding <=> :embedding) > :matchImageThreshold')
+      .orderBy('similarity', 'DESC')
+      .offset(offset)
+      .limit(limit)
+      .setParameters({
+        chatId,
+        embedding: textEmbeddingString,
+        matchImageThreshold: this.configService.get('MATCH_TEXT_THRESHOLD'),
+      })
+      .getRawMany<Messages>();
+    const t2 = performance.now();
+    console.log(`DB query time: ${Math.round(t2 - t1)} ms`);
+    // When similar
+    if (messages.length > 0) {
+      const hasMore = messages.length === limit;
+      if (firstMessageId) {
         await ctx.reply('🔎 Ось, що мені вдалось знайти:', {
-          reply_parameters: { message_id: messageId },
+          reply_parameters: { message_id: firstMessageId },
         });
-        for (const { messageId, similarity } of messages) {
-          try {
-            await ctx.reply(`${ctx.payload} (${similarity.toPrecision(4)})`, {
-              reply_parameters: { message_id: Number(messageId), allow_sending_without_reply: true },
-              disable_notification: true,
-            });
-          } catch (e) {
-            console.log(`messageId: ${messageId}`, e);
-          }
-          // Wait 300ms before send next message
-          await new Promise((r) => setTimeout(r, 300));
+      }
+      for (const message of messages) {
+        const { messageId, similarity } = message;
+        const isLast = message === messages.at(-1);
+        let reply_markup: InlineKeyboardMarkup | undefined;
+        if (isLast && hasMore) {
+          const payload = JSON.stringify({ t: text, o: offset + limit });
+          reply_markup = { inline_keyboard: [[{ text: 'Ще', callback_data: `islm-${payload}` }]] };
         }
-      } else {
+        try {
+          await ctx.reply(`${text} (${similarity.toPrecision(4)})`, {
+            reply_parameters: { message_id: Number(messageId), allow_sending_without_reply: true },
+            disable_notification: true,
+            reply_markup,
+          });
+        } catch (e) {
+          console.log(`messageId: ${messageId}`, e);
+        }
+        // Wait 1 second before send next message
+        await new Promise((r) => setTimeout(r, 1000));
+      }
+      if (!hasMore) {
+        await ctx.reply('💃 Це все!');
+      }
+    } else {
+      if (firstMessageId) {
         await ctx.reply('🤷‍♂️ Нічого нема.', {
-          reply_parameters: { message_id: messageId },
+          reply_parameters: { message_id: firstMessageId },
         });
+      } else {
+        await ctx.reply('💃 Це все!');
       }
-    } else {
-      await ctx.reply(`ℹ️ Додай пошуковий запит після команди, наприклад: /${this.command} ігрова консоль`, {
-        reply_parameters: { message_id: ctx.message.message_id },
-      });
     }
   }