Zain-ul-din · Zain-ul-din · Sep 6, 2024 · Aug 24, 2024 · Aug 24, 2024 · Aug 24, 2024
diff --git a/.gitignore b/.gitignore
@@ -16,3 +16,6 @@
 # local session
 /session
 /auth_info_baileys
+
+
+package-lock.json
diff --git a/package.json b/package.json
@@ -19,13 +19,14 @@
         "vite-node": "^0.28.5"
     },
     "dependencies": {
-        "@google/generative-ai": "^0.1.2",
+        "@google/generative-ai": "^0.17.1",
         "@types/qrcode-terminal": "^0.12.0",
         "@whiskeysockets/baileys": "^6.7.7",
         "dotenv": "^16.0.3",
         "openai": "^4.56.0",
         "ora": "3.2.0",
         "qrcode-terminal": "^0.12.0",
-        "stability-client": "^1.8.0"
+        "stability-client": "^1.8.0",
+        "whatsapp-ai-bot": "file:"
     }
 }
diff --git a/src/lib/baileys/handlers/message.ts b/src/lib/baileys/handlers/message.ts
@@ -1,23 +1,91 @@
+import { downloadMediaMessage } from '@whiskeysockets/baileys';
 import { ChatGptModel } from '../../../models/ChatGptModel';
+import { GeminiModel } from '../../../models/GeminiModel';
+import { GeminiVisionModel } from '../../../models/GeminiVisionModel';
 import { AiModels } from '../../../types/AiModels';
 import { Util } from '../../../util/Util';
 import MessageHandlerParams from '../types';
+import { DalleModel } from '../../../models/DalleModel';
 
 const gptModel = new ChatGptModel();
+const geminiModel = new GeminiModel();
+const geminiVisionModel = new GeminiVisionModel();
+const dalleModel = new DalleModel();
 
 // handles message
 export async function handleMessage({ client, msg, metadata }: MessageHandlerParams) {
   const modelToUse = Util.getModelByPrefix(metadata.text) as AiModels;
 
   if (!modelToUse) return;
 
+  const prompt = metadata.text.split(' ').slice(1).join(' ');
+
   if (modelToUse === 'ChatGPT' && metadata.msgType === 'text') {
-    gptModel.sendMessage({ sender: metadata.sender, prompt: metadata.text }, async (res, err) => {
+    await gptModel.sendMessage({ sender: metadata.sender, prompt }, async (res, err) => {
+      client.sendMessage(metadata.remoteJid, { text: err || res }, { quoted: msg });
+    });
+    return;
+  }
+
+  if (modelToUse === 'Gemini' && metadata.msgType === 'text') {
+    await geminiModel.sendMessage({ sender: metadata.sender, prompt }, async (res, err) => {
       client.sendMessage(metadata.remoteJid, { text: err || res }, { quoted: msg });
     });
     return;
   }
 
+  if (modelToUse === 'DALLE' && metadata.msgType === 'text') {
+    await dalleModel.sendMessage({ sender: metadata.sender, prompt }, async (res, err) => {
+      if (!err) {
+        client.sendMessage(
+          metadata.remoteJid,
+          { image: { url: res.url }, caption: res.caption },
+          { quoted: msg }
+        );
+      } else {
+        client.sendMessage(metadata.remoteJid, { text: err }, { quoted: msg });
+      }
+    });
+    return;
+  }
+
+  if (modelToUse === 'DALLE3' && metadata.msgType === 'text') {
+    await dalleModel.sendMessage(
+      { sender: metadata.sender, prompt, model: 'dall-e-3' },
+      async (res, err) => {
+        if (!err) {
+          client.sendMessage(
+            metadata.remoteJid,
+            { image: { url: res.url }, caption: res.caption },
+            { quoted: msg }
+          );
+        } else {
+          client.sendMessage(metadata.remoteJid, { text: err }, { quoted: msg });
+        }
+      }
+    );
+    return;
+  }
+
+  // Generative Models
+  if (
+    modelToUse === 'GeminiVision' &&
+    metadata.msgType === 'image' &&
+    metadata.imgMetaData.caption
+  ) {
+    const { mimeType, url } = metadata.imgMetaData;
+    if (mimeType === 'image/jpeg') {
+      const buffer = await downloadMediaMessage(msg, 'buffer', {});
+      await geminiVisionModel.sendMessage(
+        { sender: metadata.sender, prompt: { prompt, buffer, mimeType } },
+        async (res, err) => {
+          client.sendMessage(metadata.remoteJid, { text: err || res }, { quoted: msg });
+        }
+      );
+      return;
+    }
+  }
+
   client.sendMessage(metadata.remoteJid, { text: `Hi It's WAAI BOT here 🤖` }, { quoted: msg });
 }
 

diff --git a/src/lib/baileys/hooks/useMessageParser.ts b/src/lib/baileys/hooks/useMessageParser.ts
@@ -34,7 +34,8 @@ export default async function useMessageParser(
     | 'location'
     | 'audio' = 'unknown';
   if (conversation) msgType = 'text';
-  else if (extendedTextMessage) msgType = 'extendedText';
+  else if (extendedTextMessage) msgType = 'text';
+  // msgType = 'extendedText';
   else if (imageMessage) msgType = 'image';
   else if (videoMessage) msgType = 'video';
   else if (audioMessage) msgType = 'audio';
@@ -51,7 +52,9 @@ export default async function useMessageParser(
     msgType,
     type,
     isQuoted: false,
-    timeStamp: new Date(message.messageTimestamp * 1000),
+    timeStamp: message.messageTimestamp
+      ? new Date((message.messageTimestamp as number) * 1000)
+      : new Date(),
     text: extendedTextMessage?.text || conversation || '',
     isGroup: false,
     groupMetaData: {
@@ -74,7 +77,10 @@ export default async function useMessageParser(
     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;
+    if (imageMessage.caption) {
+      metaData.imgMetaData.caption = imageMessage.caption;
+      metaData.text = imageMessage.caption;
+    }
   }
 
   // Handle audio messages

diff --git a/src/models/ChatGptModel.ts b/src/models/ChatGptModel.ts
@@ -12,9 +12,9 @@ interface ChatGptModelParams {
   prompt: string;
 }
 
-type HandlerType = (res: string, error?: string) => Promise<void>;
+type HandleType = (res: string, error?: string) => Promise<void>;
 
-class ChatGptModel extends AiModel<ChatGptModelParams, HandlerType> {
+class ChatGptModel extends AiModel<ChatGptModelParams, HandleType> {
   public constructor() {
     super(ENV.openAIKey, 'ChatGPT');
     this.openai = new OpenAI({
@@ -26,7 +26,7 @@ class ChatGptModel extends AiModel<ChatGptModelParams, HandlerType> {
 
   public async sendMessage(
     { prompt, 
9E7A
sender }: ChatGptModelParams,
-    handler: HandlerType
+    handle: HandleType
   ): Promise<void> {
     const spinner = useSpinner(MessageTemplates.requestStr(this.aiModelName, sender, prompt));
     spinner.start();
@@ -47,7 +47,7 @@ class ChatGptModel extends AiModel<ChatGptModelParams, HandlerType> {
       this.history[sender].push(message);
 
       const res = message.content || '';
-      await handler(res);
+      await handle(res);
 
       spinner.succeed(
         MessageTemplates.reqSucceedStr(this.aiModelName, sender, res, Date.now() - startTime)
@@ -60,7 +60,7 @@ class ChatGptModel extends AiModel<ChatGptModelParams, HandlerType> {
           err
         )
       );
-      await handler('', 'An error occur please see console for more information.');
+      await handle('', 'An error occur please see console for more information.');
     }
   }
 

diff --git a/src/models/DalleModel.ts b/src/models/DalleModel.ts
@@ -1,55 +1,66 @@
-import { Message } from 'whatsapp-web.js';
 import { AiModel } from './AiModel';
 import { useSpinner } from '../hooks/useSpinner';
 import OpenAI from 'openai';
-import { MessageMedia } from 'whatsapp-web.js';
 
 import { MessageTemplates } from '../util/MessageTemplates';
 import { ENV } from '../lib/env';
 
-class DalleModel extends AiModel<string> {
-    public constructor() {
-        super(ENV.openAIKey, 'DALLE');
-        this.client = new OpenAI({ apiKey: this.apiKey });
-    }
+interface DalleModelParams {
+      sender: string;
+  prompt: string;
+  model?: 'dall-e-2' | 'dall-e-3';
+}
+
+type HandleType = (
+  { url, caption }: { url: string; caption?: string },
+  error?: string
+) => Promise<void>;
+
+class DalleModel extends AiModel<DalleModelParams, {}> {
+  public constructor() {
+    super(ENV.openAIKey, 'DALLE');
+    this.client = new OpenAI({ apiKey: this.apiKey });
+  }
+
+  public async sendMessage(
+    { sender, prompt, model }: DalleModelParams,
+    handle: HandleType
+  ): Promise<void> {
+    const spinner = useSpinner(MessageTemplates.requestStr(this.aiModelName, sender, prompt));
+    spinner.start();
+    try {
+      const startTime = Date.now();
+
+      const res = await this.client.images.generate({
+        model: model || 'dall-e-2',
+        n: 1,
+        size: '512x512',
+        prompt
+      });
+
+      await handle({ url: res.data[0].url as string, caption: res.data[0].revised_prompt });
 
-    public async sendMessage(prompt: string, msg: Message): Promise<void> {
-        const spinner = useSpinner(MessageTemplates.requestStr(this.aiModelName, msg.from, prompt));
-        spinner.start();
-        try {
-            const startTime = Date.now();
-            const response = await this.client.createImage({
-                prompt: prompt,
-                n: 1,
-                size: '512x512',
-                response_format: 'b64_json'
-            });
-
-            const base64 = response.data.data[0].b64_json as string;
-            const image = new MessageMedia('image/jpeg', base64, 'image.jpg');
-
-            msg.reply(image);
-            spinner.succeed(
-                MessageTemplates.reqSucceedStr(
-                    this.aiModelName,
-                    msg.from,
-                    response.data.data[0].url as string,
-                    Date.now() - startTime
-                )
-            );
-        } catch (err) {
-            spinner.fail(
-                MessageTemplates.reqFailStr(
-                    this.aiModelName,
-                    'at DalleModel.ts sendMessage(prompt, msg)',
-                    err
-                )
-            );
-            msg.reply('An error occur please see console for more information.');
-        }
+      spinner.succeed(
+        MessageTemplates.reqSucceedStr(
+          this.aiModelName,
+          sender,
+          res.data[0].url as string,
+          Date.now() - startTime
+        )
+      );
+    } catch (err) {
+      spinner.fail(
+        MessageTemplates.reqFailStr(
+          this.aiModelName,
+          'at DalleModel.ts sendMessage(prompt, handle)',
+          err
+        )
+      );
+      await handle({ url: '' }, 'An error occur please see console for more information.');
     }
+  }
 
-    private client;
+  private client: OpenAI;
 }
 
 export { DalleModel };