From c00ea5de983e141ef51468abdc5a8227c5caed0f Mon Sep 17 00:00:00 2001 From: Kevin Dang <77701718+kevinthedang@users.noreply.github.com> Date: Thu, 31 Jul 2025 19:04:51 -0700 Subject: [PATCH] Additional Channel Awareness (#186) --- src/client.ts | 5 +- src/events/messageCreate.ts | 68 ++++++++++++++++++++---- src/utils/events.ts | 4 +- src/utils/handlers/chatHistoryHandler.ts | 33 ++++++++++++ 4 files changed, 97 insertions(+), 13 deletions(-) diff --git a/src/client.ts b/src/client.ts index 82d76ca..fbf72f5 100644 --- a/src/client.ts +++ b/src/client.ts @@ -23,8 +23,11 @@ export const ollama = new Ollama({ // Create Queue managed by Events const messageHistory: Queue = new Queue +// Create Channel History Queue managed by Events +const channelMessageHistory: Queue = new Queue + // register all events -registerEvents(client, Events, messageHistory, ollama, Keys.defaultModel) +registerEvents(client, Events, messageHistory, channelMessageHistory, ollama, Keys.defaultModel) // Try to log in the client await client.login(Keys.clientToken) diff --git a/src/events/messageCreate.ts b/src/events/messageCreate.ts index 587dc48..26b1f29 100644 --- a/src/events/messageCreate.ts +++ b/src/events/messageCreate.ts @@ -1,5 +1,5 @@ import { TextChannel } from 'discord.js' -import { event, Events, normalMessage, UserMessage, clean } from '../utils/index.js' +import { event, Events, normalMessage, UserMessage, clean, addToChannelContext } from '../utils/index.js' import { getChannelInfo, getServerConfig, getUserConfig, openChannelInfo, openConfig, UserConfig, getAttachmentData, getTextFileAttachmentData @@ -11,7 +11,7 @@ import { * * @param message the message received from the channel */ -export default event(Events.MessageCreate, async ({ log, msgHist, ollama, client, defaultModel }, message) => { +export default event(Events.MessageCreate, async ({ log, msgHist, channelHistory, ollama, client, defaultModel }, message) => { const clientId = client.user!!.id let cleanedMessage = clean(message.content, clientId) log(`Message \"${cleanedMessage}\" from ${message.author.tag} in channel/thread ${message.channelId}.`) @@ -19,6 +19,61 @@ export default event(Events.MessageCreate, async ({ log, msgHist, ollama, client // Do not respond if bot talks in the chat if (message.author.username === message.client.user.username) return + // Save User Chat even if not for the bot + let channelContextHistory: UserMessage[] = await new Promise((resolve) => { + getChannelInfo(`${message.channelId}-context.json`, (channelInfo) => { + if (channelInfo?.messages) + resolve(channelInfo.messages) + else { + log(`Channel/Thread ${message.channel}-context does not exist. File will be created shortly...`) + resolve([]) + } + }) + }) + + if (channelContextHistory.length === 0) { + channelContextHistory = await new Promise((resolve) => { + addToChannelContext(message.channelId, + message.channel as TextChannel + ) + getChannelInfo(`${message.channelId}-context.json`, (channelInfo) => { + if (channelInfo?.messages) + resolve(channelInfo.messages) + else { + log(`Channel/Thread ${message.channel}-context does not exist. File will be created shortly...`) + } + }) + }) + } + + // Set Channel History Queue + channelHistory.setQueue(channelContextHistory) + + // get message attachment if exists + const attachment = message.attachments.first() + let messageAttachment: string[] = [] + + if (attachment && attachment.name?.endsWith(".txt")) + cleanedMessage += ' ' + await getTextFileAttachmentData(attachment) + else if (attachment) + messageAttachment = await getAttachmentData(attachment) + + while (channelHistory.size() >= channelHistory.capacity) channelHistory.dequeue() + + // push user response to channel history + console.log + channelHistory.enqueue({ + role: 'user', + content: cleanedMessage, + images: messageAttachment || [] + }) + + // Store in Channel Context + addToChannelContext(message.channelId, + message.channel as TextChannel, + channelHistory.getItems() + ) + // Only respond if message mentions the bot if (!message.mentions.has(clientId)) return @@ -139,15 +194,6 @@ export default event(Events.MessageCreate, async ({ log, msgHist, ollama, client if (!userConfig) throw new Error(`Failed to initialize User Preference for **${message.author.username}**.\n\nIt's likely you do not have a model set. Please use the \`switch-model\` command to do that.`) - // get message attachment if exists - const attachment = message.attachments.first() - let messageAttachment: string[] = [] - - if (attachment && attachment.name?.endsWith(".txt")) - cleanedMessage += await getTextFileAttachmentData(attachment) - else if (attachment) - messageAttachment = await getAttachmentData(attachment) - const model: string = userConfig.options['switch-model'] // set up new queue diff --git a/src/utils/events.ts b/src/utils/events.ts index 6bb34a4..f6939a4 100644 --- a/src/utils/events.ts +++ b/src/utils/events.ts @@ -37,6 +37,7 @@ export interface EventProps { client: Client, log: LogMethod, msgHist: Queue, + channelHistory: Queue, ollama: Ollama, defaultModel: String } @@ -78,6 +79,7 @@ export function registerEvents( client: Client, events: Event[], msgHist: Queue, + channelHistory: Queue, ollama: Ollama, defaultModel: String ): void { @@ -88,7 +90,7 @@ export function registerEvents( // Handle Errors, call callback, log errors as needed try { - callback({ client, log, msgHist, ollama, defaultModel }, ...args) + callback({ client, log, msgHist, channelHistory, ollama, defaultModel }, ...args) } catch (error) { log('[Uncaught Error]', error) } diff --git a/src/utils/handlers/chatHistoryHandler.ts b/src/utils/handlers/chatHistoryHandler.ts index 186e342..9bcf403 100644 --- a/src/utils/handlers/chatHistoryHandler.ts +++ b/src/utils/handlers/chatHistoryHandler.ts @@ -56,6 +56,39 @@ export async function clearChannelInfo(filename: string, channel: TextChannel, u return cleanedHistory } +export async function addToChannelContext(filename: string, channel : TextChannel | ThreadChannel, messages: UserMessage[] = []): Promise { + const fullFileName = `data/${filename}-context.json` + if (fs.existsSync(fullFileName)) { + fs.readFile(fullFileName, 'utf8', (error, data) => { + if (error) + console.log(`[Error: addToChannelContext] Incorrect file format`) + else { + const object = JSON.parse(data) + if (object['messages'].length === 0) + object['messages'] = messages as [] + else if (object['messages'].length !== 0 && messages.length !== 0) + object['messages'] = messages as [] + fs.writeFileSync(fullFileName, JSON.stringify(object, null, 2)) + } + }) + } else { // channel context does not exist, create it + const object: Configuration = JSON.parse( + `{ + \"id\": \"${channel?.id}\", + \"name\": \"${channel?.name}\", + \"messages\": [] + }` + ) + + const directory = path.dirname(fullFileName) + if (!fs.existsSync(directory)) + fs.mkdirSync(directory, { recursive: true }) + + fs.writeFileSync(fullFileName, JSON.stringify(object, null, 2)) + console.log(`[Util: addToChannelContext] Created '${fullFileName}' in working directory`) + } +} + /** * Method to open the channel history *