diff --git a/docker-compose.yml b/docker-compose.yml index f404cd6..2fc820d 100644 --- a/docker-compose.yml +++ b/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} diff --git a/package-lock.json b/package-lock.json index aa8707f..b9ed039 100644 --- a/package-lock.json +++ b/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", diff --git a/package.json b/package.json index f31766f..1e2f2d9 100644 --- a/package.json +++ b/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", diff --git a/src/client.ts b/src/client.ts index f0d4dc4..0494674 100644 --- a/src/client.ts +++ b/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: [] }) \ No newline at end of file diff --git a/src/events/messageCreate.ts b/src/events/messageCreate.ts index 00d3ff2..39dd558 100644 --- a/src/events/messageCreate.ts +++ b/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 diff --git a/src/utils/events.ts b/src/utils/events.ts index 082e78a..ae83074 100644 --- a/src/utils/events.ts +++ b/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 diff --git a/src/utils/handlers/bufferHandler.ts b/src/utils/handlers/bufferHandler.ts new file mode 100644 index 0000000..d0ef4e1 --- /dev/null +++ b/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 { + // 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 { + 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] +} \ No newline at end of file diff --git a/src/utils/index.ts b/src/utils/index.ts index edae3a1..4e3e5ea 100644 --- a/src/utils/index.ts +++ b/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'