Skip to content

Commit

Permalink
Start import media history by command
Browse files Browse the repository at this point in the history
  • Loading branch information
@OperKH
OperKH committed Feb 26, 2024
1 parent 5fbe746 commit 69f9764
Show file tree
Hide file tree
Showing 3 changed files with 135 additions and 99 deletions.
92 changes: 46 additions & 46 deletions package-lock.json
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 3 additions & 3 deletions package.json
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "ai_bot",
"version": "1.2.0",
"version": "1.2.1",
"private": true,
"main": "src/app.ts",
"type": "module",
Expand Down Expand Up @@ -30,8 +30,8 @@
"devDependencies": {
"@types/fluent-ffmpeg": "^2.1.24",
"@types/node": "^20.11.20",
"@typescript-eslint/eslint-plugin": "^7.0.2",
"@typescript-eslint/parser": "^7.0.2",
"@typescript-eslint/eslint-plugin": "^7.1.0",
"@typescript-eslint/parser": "^7.1.0",
"eslint": "^8.57.0",
"nodemon": "^3.1.0",
"prettier": "^3.2.5",
Expand Down
136 changes: 86 additions & 50 deletions src/bot/commands/mediaTracker.command.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ export class MediaTrackerCommand extends Command {
private fileService = FileService.getInstance();
private tgClient: TelegramClient | null = null;
private similarFoundVariants = ['ось тут', 'ще тут', 'і ось', 'навіть це', 'і оце щось схоже'];
private isMediaImporting = false;

handle(): void {
this.bot.on(message('photo'), async (ctx) => {
Expand All @@ -35,38 +36,23 @@ export class MediaTrackerCommand extends Command {
reply_parameters: { message_id: ctx.message.message_id },
});
});
this.bot.command('starthistoryimport', async (ctx) => {
this.startHistoryImport(ctx);
});
}

private async messageHandler(ctx: NarrowedContext<IBotContext, Update.MessageUpdate<Message>>, fileId: string) {
if (this.isMediaImporting) return;
const chatId = ctx.chat.id;
const messageId = ctx.message.message_id;
const chatStateRepository = this.dataSource.getRepository(ChatState);
const chatPhotoMessageRepository = this.dataSource.getRepository(ChatPhotoMessage);

try {
// Check media import is done
const chatState = await chatStateRepository.findOneBy({ chatId: String(chatId) });
const isMediaImported = chatState?.isMediaImported ?? false;
if (!isMediaImported) {
const [latestChatPhotoMessage] = await chatPhotoMessageRepository.find({
where: { chatId: String(chatId) },
order: { messageId: 'DESC' },
take: 1,
});
const lastImportedMessageId = latestChatPhotoMessage ? Number(latestChatPhotoMessage.messageId) : 0;
const lastMessageId = messageId;
await this.importChatMessages(chatId, lastImportedMessageId, lastMessageId);
// Save to DB
const chatState = new ChatState();
chatState.chatId = String(chatId);
(chatState.isMediaImported = true), await chatStateRepository.save(chatState);
}
// DB similarity search
const imageEmbeddingString = await this.getEmbeddingStringByImageFileId(fileId);
type Messages = {
messageId: string;
similarity: number;
};
const chatPhotoMessageRepository = this.dataSource.getRepository(ChatPhotoMessage);
const messages = await chatPhotoMessageRepository
.createQueryBuilder('msg')
.select('msg.messageId', 'messageId')
Expand Down Expand Up @@ -107,6 +93,56 @@ export class MediaTrackerCommand extends Command {
}
}

private async startHistoryImport(ctx: NarrowedContext<IBotContext, Update.MessageUpdate<Message>>) {
const messageId = ctx.message.message_id;
if (this.isMediaImporting) {
await ctx.reply('😡 Я тут працюю, тужуся, а ти відволікаєш', {
reply_parameters: { message_id: messageId },
});
return;
}
this.isMediaImporting = true;
try {
const chatId = ctx.chat.id;
const chatStateRepository = this.dataSource.getRepository(ChatState);
const chatState = await chatStateRepository.findOneBy({ chatId: String(chatId) });
const isMediaImported = chatState?.isMediaImported ?? false;
if (isMediaImported) {
await ctx.reply('🍧 Нема потреби. Усе же зроблено.', {
reply_parameters: { message_id: messageId },
});
} else {
await ctx.reply('🏃 Взяв у роботу!', {
reply_parameters: { message_id: messageId },
});

const chatPhotoMessageRepository = this.dataSource.getRepository(ChatPhotoMessage);
const [latestChatPhotoMessage] = await chatPhotoMessageRepository.find({
where: { chatId: String(chatId) },
order: { messageId: 'DESC' },
take: 1,
});
const lastImportedMessageId = latestChatPhotoMessage ? Number(latestChatPhotoMessage.messageId) : 0;
await this.importChatMessages(chatId, lastImportedMessageId, messageId);
// Save to DB
const chatState = new ChatState();
chatState.chatId = String(chatId);
(chatState.isMediaImported = true), await chatStateRepository.save(chatState);

await ctx.reply('😮‍💨 Фух... Усе підтягнув!', {
reply_parameters: { message_id: messageId },
});
}
} catch (e) {
console.log(e);
await ctx.reply('📛 Халепа', {
reply_parameters: { message_id: messageId },
});
} finally {
this.isMediaImporting = false;
}
}

private async getEmbeddingStringByImageFileId(fileId: string) {
const fileUrl = await this.bot.telegram.getFileLink(fileId);
const imageBuffer = await this.fileService.getBufferByUrl(fileUrl);
Expand All @@ -132,36 +168,36 @@ export class MediaTrackerCommand extends Command {
reverse: true,
filter: new Api.InputMessagesFilterPhotoVideo(),
})) {
if (message.id < lastMessageId) {
const t1 = performance.now();
let fleLocation: Api.InputPhotoFileLocation | Api.InputDocumentFileLocation | null = null;
if (message.video) {
const { id, fileReference, accessHash } = message.video;
const thumbSize = message.video.thumbs?.findLast(({ className }) => className === 'PhotoSize')?.type ?? 'm';
fleLocation = new Api.InputDocumentFileLocation({ id, fileReference, accessHash, thumbSize });
} else if (message.photo) {
const photo = message.photo as Api.Photo;
const { id, fileReference, accessHash } = photo;
const thumbSize = photo.sizes.at(-1)?.type ?? 'm';
fleLocation = new Api.InputPhotoFileLocation({ id, fileReference, accessHash, thumbSize });
}
if (fleLocation) {
const imageBuffer = await this.tgClient.downloadFile(fleLocation);
if (imageBuffer instanceof Buffer) {
// Get image embedding
const rawImage = await this.aiService.getRawImageFromBuffer(imageBuffer);
const imageEmbedding = await this.aiService.getImageClipEmbedding(rawImage);
const imageEmbeddingString = JSON.stringify(imageEmbedding);
// Save to DB
const chatPhotoMessage = new ChatPhotoMessage();
chatPhotoMessage.chatId = String(chatId);
chatPhotoMessage.messageId = String(message.id);
chatPhotoMessage.embedding = imageEmbeddingString;
await this.dataSource.manager.save(chatPhotoMessage);
// Logging
const t2 = performance.now();
console.log(`Imported message ${message.id}/${lastMessageId - 1} (${Math.round(t2 - t1)} ms)`);
}
const t1 = performance.now();
let fleLocation: Api.InputPhotoFileLocation | Api.InputDocumentFileLocation | null = null;
if (message.video) {
const { id, fileReference, accessHash } = message.video;
const thumbSize = message.video.thumbs?.findLast(({ className }) => className === 'PhotoSize')?.type ?? 'm';
fleLocation = new Api.InputDocumentFileLocation({ id, fileReference, accessHash, thumbSize });
} else if (message.photo) {
const photo = message.photo as Api.Photo;
const { id, fileReference, accessHash } = photo;
const thumbSize = photo.sizes.at(-1)?.type ?? 'm';
fleLocation = new Api.InputPhotoFileLocation({ id, fileReference, accessHash, thumbSize });
}
if (fleLocation) {
const imageBuffer = await this.tgClient.downloadFile(fleLocation);
if (imageBuffer instanceof Buffer) {
// Get image embedding
const rawImage = await this.aiService.getRawImageFromBuffer(imageBuffer);
const imageEmbedding = await this.aiService.getImageClipEmbedding(rawImage);
const imageEmbeddingString = JSON.stringify(imageEmbedding);
// Save to DB
const chatPhotoMessage = new ChatPhotoMessage();
chatPhotoMessage.chatId = String(chatId);
chatPhotoMessage.messageId = String(message.id);
chatPhotoMessage.embedding = imageEmbeddingString;
await this.dataSource.manager.save(chatPhotoMessage);
// Logging
const t2 = performance.now();
console.log(
`Imported message ${message.id}/${lastMessageId} ${Math.round((message.id / lastMessageId) * 1e4) / 1e2}% (${Math.round(t2 - t1)} ms)`,
);
}
}
}
Expand Down

0 comments on commit 69f9764

Please sign in to comment.