From c234cdd07211c1b13c77a8830f93ed926982ef39 Mon Sep 17 00:00:00 2001 From: Vicente Date: Sat, 23 Nov 2024 13:54:46 -0300 Subject: [PATCH 1/3] Added IGNORE_SELF_MESSAGES configuration and removed self message handle --- .env.example | 1 + src/baileys/env.ts | 2 ++ src/baileys/handlers/message.ts | 9 +-------- src/baileys/handlers/messages.ts | 8 ++++---- 4 files changed, 8 insertions(+), 12 deletions(-) diff --git a/.env.example b/.env.example index aee5c91..9d5bb89 100644 --- a/.env.example +++ b/.env.example @@ -5,6 +5,7 @@ DEBUG=False PROCESSING="Thinking..." +IGNORE_SELF_MESSAGES=False # Model Services API_KEY_OPENAI=ADD_YOUR_KEY diff --git a/src/baileys/env.ts b/src/baileys/env.ts index 8b5e272..d239460 100644 --- a/src/baileys/env.ts +++ b/src/baileys/env.ts @@ -6,6 +6,7 @@ configEnv(); interface EnvInterface { Debug: boolean; Processing: string; + IGNORE_SELF_MESSAGES: boolean; // Model Services API_KEY_OPENAI?: string; @@ -52,6 +53,7 @@ interface EnvInterface { export const ENV: EnvInterface = { Debug: process.env.DEBUG === 'True', Processing: process.env.PROCESSING || 'Processing...', + IGNORE_SELF_MESSAGES: process.env.IGNORE_SELF_MESSAGES === 'True', API_KEY_OPENAI: process.env.API_KEY_OPENAI, API_KEY_OPENAI_DALLE: process.env.API_KEY_OPENAI_DALLE || process.env.API_KEY_OPENAI, diff --git a/src/baileys/handlers/message.ts b/src/baileys/handlers/message.ts index 73b5a8b..8ddac55 100644 --- a/src/baileys/handlers/message.ts +++ b/src/baileys/handlers/message.ts @@ -100,11 +100,4 @@ export async function handleMessage({ client, msg, metadata }: MessageHandlerPar } } ); -} - -// handles message from self -export async function handlerMessageFromMe({ metadata, client, msg, type }: MessageHandlerParams) { - // if (metadata.fromMe && metadata.isQuoted) return; - // if (metadata.isQuoted && Util.getModelByPrefix(metadata.text)) return; - await handleMessage({ metadata, client, msg, type }); -} +} \ No newline at end of file diff --git a/src/baileys/handlers/messages.ts b/src/baileys/handlers/messages.ts index 2f5cb4e..4203283 100644 --- a/src/baileys/handlers/messages.ts +++ b/src/baileys/handlers/messages.ts @@ -1,6 +1,7 @@ import makeWASocket, { MessageUpsertType, WAMessage } from '@whiskeysockets/baileys'; import useMessageParser from '../hooks/useMessageParser'; -import { handleMessage, handlerMessageFromMe } from './message'; +import { handleMessage } from './message'; +import { ENV } from '../env'; export function messagesHandler({ client, @@ -21,10 +22,9 @@ export function messagesHandler({ if (!metadata) return; if (metadata.msgType === 'unknown') return; if (metadata.isGroup && metadata.groupMetaData.groupIsLocked) return; + if (metadata.fromMe && ENV.IGNORE_SELF_MESSAGES) return; - await (metadata.fromMe - ? handlerMessageFromMe({ client, msg, metadata, type }) - : handleMessage({ client, msg, metadata, type })); + await handleMessage({ client, msg, metadata, type }); } catch (_) {} }) ); From 12d5f158f3cf924bb3247a7335164afe03b26a4a Mon Sep 17 00:00:00 2001 From: Vicente Date: Sat, 23 Nov 2024 13:55:04 -0300 Subject: [PATCH 2/3] Enhance message parser metadata to include hasImage and hasAudio flags --- src/baileys/hooks/useMessageParser.ts | 25 +++++++++++++++---------- src/models/BaseAiModel.ts | 4 ++++ 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/src/baileys/hooks/useMessageParser.ts b/src/baileys/hooks/useMessageParser.ts index bff713d..3c08278 100644 --- a/src/baileys/hooks/useMessageParser.ts +++ b/src/baileys/hooks/useMessageParser.ts @@ -49,6 +49,7 @@ export default async function useMessageParser( senderName: message.pushName, fromMe, msgType, + message, type, isQuoted: false, quoteMetaData: { @@ -56,6 +57,7 @@ export default async function useMessageParser( message: {}, text: '', type: 'text', + hasImage: false, imgMetaData: { url: '', mimeType: '', @@ -71,11 +73,13 @@ export default async function useMessageParser( groupName: '', groupIsLocked: false }, + hasImage: false, imgMetaData: { url: '', mimeType: '', caption: '' }, + hasAudio: false, audioMetaData: { url: '', mimeType: '' @@ -84,20 +88,18 @@ export default async function useMessageParser( // Handle image messages if (imageMessage) { - metaData.msgType = 'image'; - if (imageMessage.url) metaData.imgMetaData.url = imageMessage.url; - if (imageMessage.mimetype) metaData.imgMetaData.mimeType = imageMessage.mimetype; - if (imageMessage.caption) { - metaData.imgMetaData.caption = imageMessage.caption; - metaData.text = imageMessage.caption; - } + metaData.hasImage = true; + metaData.imgMetaData.url = imageMessage.url || ''; + metaData.imgMetaData.mimeType = imageMessage.mimetype || ''; + metaData.imgMetaData.caption = imageMessage.caption || ''; + metaData.text = imageMessage.caption || ''; } // Handle audio messages if (audioMessage) { - metaData.msgType = 'audio'; - if (audioMessage.url) metaData.audioMetaData.url = audioMessage.url; - if (audioMessage.mimetype) metaData.audioMetaData.mimeType = audioMessage.mimetype; + metaData.hasAudio = true; + metaData.audioMetaData.url = audioMessage.url || ''; + metaData.audioMetaData.mimeType = audioMessage.mimetype || ''; } // gather context info @@ -112,6 +114,7 @@ export default async function useMessageParser( metaData.quoteMetaData.message = contextInfo.quotedMessage || {}; if (contextInfo.quotedMessage.imageMessage) { + metaData.quoteMetaData.hasImage = true; metaData.quoteMetaData.type = 'image'; metaData.quoteMetaData.imgMetaData.url = contextInfo.quotedMessage.imageMessage.url || ''; metaData.quoteMetaData.imgMetaData.mimeType = @@ -129,5 +132,7 @@ export default async function useMessageParser( metaData.groupMetaData.groupIsLocked = groupMetadata.restrict || false; } + metaData.message = message; + return metaData; } diff --git a/src/models/BaseAiModel.ts b/src/models/BaseAiModel.ts index ce15ad3..a230050 100644 --- a/src/models/BaseAiModel.ts +++ b/src/models/BaseAiModel.ts @@ -8,6 +8,7 @@ type AIMetaData = { sender: string; senderName?: string; fromMe: boolean; + message: any; msgType: | 'unknown' | 'text' @@ -25,6 +26,7 @@ type AIMetaData = { message: any; text: string; type: 'text' | 'image'; + hasImage: boolean; imgMetaData: { url: string; mimeType: string; @@ -38,11 +40,13 @@ type AIMetaData = { groupName: string; groupIsLocked: boolean; }; + hasImage: boolean; imgMetaData: { url: string; mimeType: string; caption: string; }; + hasAudio: boolean; audioMetaData: { url: string; mimeType: string; From 7435aca77ef9b42ae1e5b9d4d9c9433d402fd6ef Mon Sep 17 00:00:00 2001 From: Vicente Date: Sat, 23 Nov 2024 13:55:18 -0300 Subject: [PATCH 3/3] Refactor GeminiModel to improve image handling and validation in generateImageCompletion method --- src/models/GeminiModel.ts | 75 +++++++++++++++++++++++++-------------- 1 file changed, 49 insertions(+), 26 deletions(-) diff --git a/src/models/GeminiModel.ts b/src/models/GeminiModel.ts index ac7f8fb..e34c887 100644 --- a/src/models/GeminiModel.ts +++ b/src/models/GeminiModel.ts @@ -10,12 +10,20 @@ import { downloadMediaMessage } from '@whiskeysockets/baileys'; /* Local modules */ import { AIModel, AIArguments, AIHandle, AIMetaData } from './BaseAiModel'; import { ENV } from '../baileys/env'; -import invariant from 'invariant'; +// import invariant from 'invariant'; + +interface imgMetaData { + url: string; + mimeType: string; + caption: string; +} + +const validMimeTypes = ['image/jpeg', 'image/png', 'image/jpg', 'image/webp']; /* Gemini Model */ class GeminiModel extends AIModel { /* Variables */ - private generativeModel?: GenerativeModel; + private generativeModel: GenerativeModel; private Gemini: GoogleGenerativeAI; public chats: { [from: string]: ChatSession }; public instructions: string | undefined; @@ -23,6 +31,13 @@ class GeminiModel extends AIModel { public constructor() { super(ENV.API_KEY_GEMINI, 'Gemini', ENV.GEMINI_ICON_PREFIX); this.Gemini = new GoogleGenerativeAI(ENV.API_KEY_GEMINI as string); + + // https://ai.google.dev/api/generate-content#method:-models.generatecontent + this.generativeModel = this.Gemini.getGenerativeModel({ + model: 'gemini-1.5-flash', + systemInstruction: this.instructions + }); + this.chats = {}; } @@ -30,7 +45,7 @@ class GeminiModel extends AIModel { public async generateCompletion(user: string, prompt: string): Promise { if (!this.sessionExists(user)) { this.sessionCreate(user); - this.chats[user] = this.generativeModel!.startChat(); + this.chats[user] = this.generativeModel.startChat(); } const chat = this.chats[user]; @@ -46,25 +61,17 @@ class GeminiModel extends AIModel { }; } - private initGenerativeModel() { - // https://ai.google.dev/gemini-api/docs/models/gemini - this.generativeModel = this.Gemini.getGenerativeModel({ - model: 'gemini-1.5-flash', - systemInstruction: this.instructions - }); - } - - public async generateImageCompletion(prompt: string, metadata: AIMetaData): Promise { - this.initGenerativeModel(); - invariant(this.generativeModel, 'Unable to initialize Gemini Generative model'); + public async generateImageCompletion( + prompt: string, + imgMetaData: imgMetaData, + message: any + ): Promise { + // this.initGenerativeModel(); + // invariant(this.generativeModel, 'Unable to initialize Gemini Generative model'); - const { mimeType } = metadata.quoteMetaData.imgMetaData; - if (mimeType === 'image/jpeg') { - const buffer = await downloadMediaMessage( - { message: metadata.quoteMetaData.message } as any, - 'buffer', - {} - ); + const { mimeType } = imgMetaData; + if (validMimeTypes.includes(mimeType)) { + const buffer = await downloadMediaMessage({ message: message } as any, 'buffer', {}); const imageParts = this.createGenerativeContent(buffer, mimeType); const result = await this.generativeModel.generateContent([prompt, imageParts]); const resultText = result.response.text(); @@ -72,21 +79,37 @@ class GeminiModel extends AIModel { return resultText; } - return ''; + return 'The image is not a valid image type.'; } async sendMessage({ sender, prompt, metadata }: AIArguments, handle: AIHandle) { try { let message = ''; if (metadata.isQuoted) { - if (metadata.quoteMetaData.type === 'image') { - message = this.iconPrefix + (await this.generateImageCompletion(prompt, metadata)); + if (metadata.quoteMetaData.hasImage) { + message = + this.iconPrefix + + (await this.generateImageCompletion( + prompt, + metadata.quoteMetaData.imgMetaData, + metadata.quoteMetaData.message + )); } else { - prompt = 'Quoted Message:\n' + metadata.quoteMetaData.text + '---\nMessage:\n' + prompt; + prompt = 'Quoted Message:\n' + metadata.quoteMetaData.text + '\n---\nMessage:\n' + prompt; message = this.iconPrefix + (await this.generateCompletion(sender, prompt)); } } else { - message = this.iconPrefix + (await this.generateCompletion(sender, prompt)); + if (metadata.hasImage) { + message = + this.iconPrefix + + (await this.generateImageCompletion( + prompt, + metadata.imgMetaData, + metadata.message.message + )); + } else { + message = this.iconPrefix + (await this.generateCompletion(sender, prompt)); + } } handle({ text: message });