Message Attachment Image Recognition (#118)

* Add: Image recognition

* Fix: Non-Attachment Query

* Update: version increment

* Remove: Debugging logs from buffer file

* Add: comments for bufferHandler

docker-compose.yml

@@ -7,7 +7,7 @@ services:
     build: ./                     # find docker file in designated path
     container_name: discord
     restart: always               # rebuild container always
-    image: kevinthedang/discord-ollama:0.5.11
+    image: kevinthedang/discord-ollama:0.6.0
     environment:
       CLIENT_TOKEN: ${CLIENT_TOKEN}
       MODEL: ${MODEL}

package-lock.json

@@ -1,12 +1,12 @@
 {
   "name": "discord-ollama",
-  "version": "0.5.11",
+  "version": "0.6.0",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "discord-ollama",
-      "version": "0.5.11",
+      "version": "0.6.0",
       "license": "ISC",
       "dependencies": {
         "discord.js": "^14.15.3",

package.json

@@ -1,6 +1,6 @@
 {
   "name": "discord-ollama",
-  "version": "0.5.11",
+  "version": "0.6.0",
   "description": "Ollama Integration into discord",
   "main": "build/index.js",
   "exports": "./build/index.js",

src/client.ts

@@ -45,5 +45,6 @@ await client.login(Keys.clientToken)
 // queue up bots name
 messageHistory.enqueue({
     role: 'assistant',
-    content: `My name is ${client.user?.username}`
+    content: `My name is ${client.user?.username}`,
+    images: []
 })

src/events/messageCreate.ts

@@ -1,5 +1,5 @@
 import { embedMessage, event, Events, normalMessage, UserMessage } from '../utils/index.js'
-import { getChannelInfo, getServerConfig, getUserConfig, openChannelInfo, openConfig, ServerConfig, UserConfig } from '../utils/index.js'
+import { getChannelInfo, getServerConfig, getUserConfig, openChannelInfo, openConfig, UserConfig, getAttachmentData } from '../utils/index.js'
 import { clean } from '../utils/mentionClean.js'
 import { TextChannel } from 'discord.js'
 
@@ -102,6 +102,9 @@ export default event(Events.MessageCreate, async ({ log, msgHist, tokens, ollama
         // response string for ollama to put its response
         let response: string
 
+        // get message attachment if exists
+        const messageAttachment: string[] = await getAttachmentData(message.attachments.first())
+
         // set up new queue
         msgHist.setQueue(chatMessages)
 
@@ -111,7 +114,8 @@ export default event(Events.MessageCreate, async ({ log, msgHist, tokens, ollama
         // push user response before ollama query
         msgHist.enqueue({
             role: 'user',
-            content: clean(message.content)
+            content: clean(message.content),
+            images: messageAttachment || []
         })
         
         // undefined or false, use normal, otherwise use embed
@@ -129,7 +133,8 @@ export default event(Events.MessageCreate, async ({ log, msgHist, tokens, ollama
         // successful query, save it in context history
         msgHist.enqueue({ 
             role: 'assistant', 
-            content: response
+            content: response,
+            images: messageAttachment || []
         })
 
         // only update the json on success

src/utils/events.ts

@@ -38,7 +38,8 @@ export type ChatParams = {
  */
 export type UserMessage = {
     role: string,
-    content: string
+    content: string,
+    images: string[] // May or may not have images in message
 }
 
 // Event properties

src/utils/handlers/bufferHandler.ts

@@ -0,0 +1,57 @@
+import { Attachment } from "discord.js"
+
+/**
+ * Method to convert a Discord attachment url to an array buffer
+ * 
+ * @param url Discord Attachment Url
+ * @returns array buffer from Attachment Url
+ */
+async function getAttachmentBuffer(url: string): Promise<ArrayBuffer> {
+    // Get the data from the image
+    const response = await fetch(url)
+
+    // Validate the image came in fine
+    if (!response.ok)
+        throw new Error('Failed to fetch the attachment.')
+
+    // Return image as Buffer
+    return await response.arrayBuffer()
+}
+
+/**
+ * Method to convert an array buffer to a Base64 String
+ * 
+ * @param buffer Array Buffer from attachment
+ * @returns converted Base64 string
+ */
+function arrayBufferToBase64(buffer: ArrayBuffer): string {
+    // Converting to Uint8Array
+    const uint8Array = new Uint8Array(buffer)
+    let binary = ''
+    const len = uint8Array.byteLength;
+    for (let i = 0; i < len; i++) {
+        binary += String.fromCharCode(uint8Array[i])
+    }
+
+    // Return as Base64
+    return Buffer.from(binary, 'binary').toString('base64')
+}
+
+/**
+ * Method to retrieve the Base64 Array of provided Message Attachment
+ * 
+ * @param attachment Message Attachment from Discord
+ * @returns Base64 string array
+ */
+export async function getAttachmentData(attachment: Attachment | undefined): Promise<string[]> {
+    const url: string = attachment !== undefined ? attachment.url : "Missing Url"
+
+    // case of no attachment
+    if (url === "Missing Url")
+        return []
+
+    // Convert data to base64
+    const buffer = await getAttachmentBuffer(url)
+    const base64String = arrayBufferToBase64(buffer)
+    return [base64String]
+}

src/utils/index.ts

@@ -10,3 +10,4 @@ export * from './configInterfaces.js'
 export * from './handlers/chatHistoryHandler.js'
 export * from './handlers/configHandler.js'
 export * from './handlers/streamHandler.js'
+export * from './handlers/bufferHandler.js'