Streamlined Preferences Setup and Default Model (#148)

* Update: Streamlinded setup and Default Model * Update: version increment
kevinthedang · Dec 12, 2024 · 6ac45af · 6ac45af
1 parent d570a50
commit 6ac45af
Show file tree

Hide file tree

Showing 9 changed files with 100 additions and 57 deletions.
diff --git a/.env.sample b/.env.sample
@@ -1,6 +1,9 @@
 # Discord token for the bot
 CLIENT_TOKEN = BOT_TOKEN
 
+# Default model for new users
+MODEL = DEFAULT_MODEL
+
 # ip/port address of docker container, I use 172.18.0.3 for docker, 127.0.0.1 for local
 OLLAMA_IP = IP_ADDRESS
 OLLAMA_PORT = PORT

diff --git 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.7.3
+    image: kevinthedang/discord-ollama:0.7.4
     environment:
       CLIENT_TOKEN: ${CLIENT_TOKEN}
       OLLAMA_IP: ${OLLAMA_IP}

diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "discord-ollama",
-  "version": "0.7.3",
+  "version": "0.7.4",
   "description": "Ollama Integration into discord",
   "main": "build/index.js",
   "exports": "./build/index.js",

diff --git a/src/client.ts b/src/client.ts
@@ -25,7 +25,7 @@ export const ollama = new Ollama({
 const messageHistory: Queue<UserMessage> = new Queue<UserMessage>
 
 // register all events
-registerEvents(client, Events, messageHistory, ollama)
+registerEvents(client, Events, messageHistory, 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
@@ -8,7 +8,7 @@ import { getChannelInfo, getServerConfig, getUserConfig, openChannelInfo, openCo
  * 
  * @param message the message received from the channel
  */
-export default event(Events.MessageCreate, async ({ log, msgHist, ollama, client }, message) => {
+export default event(Events.MessageCreate, async ({ log, msgHist, ollama, client, defaultModel }, message) => {
     const clientId = client.user!!.id
     const cleanedMessage = clean(message.content, clientId)
     log(`Message \"${cleanedMessage}\" from ${message.author.tag} in channel/thread ${message.channelId}.`)
@@ -21,57 +21,88 @@ export default event(Events.MessageCreate, async ({ log, msgHist, ollama, client
 
     // default stream to false
     let shouldStream = false
+
+    // Params for Preferences Fetching
+    const maxRetries = 3
+    const delay = 1000 // in millisecons
 
     try {
         // Retrieve Server/Guild Preferences
-        await new Promise((resolve, reject) => {
-            getServerConfig(`${message.guildId}-config.json`, (config) => {
-                // check if config.json exists
-                if (config === undefined) {
-                    // Allowing chat options to be available
-                    openConfig(`${message.guildId}-config.json`, 'toggle-chat', true)
-                    reject(new Error('No Server Preferences is set up.\n\nCreating default server preferences file...\nPlease try chatting again.'))
-                    return
-                }
-
-                // check if chat is disabled
-                if (!config.options['toggle-chat']) {
-                    reject(new Error('Admin(s) have disabled chat features.\n\n Please contact your server\'s admin(s).'))
-                    return
-                }
-
-                resolve(config)
-            })
-        })
+        let attempt = 0
+        while (attempt < maxRetries) {
+            try {
+                await new Promise((resolve, reject) => {
+                    getServerConfig(`${message.guildId}-config.json`, (config) => {
+                        // check if config.json exists
+                        if (config === undefined) {
+                            // Allowing chat options to be available
+                            openConfig(`${message.guildId}-config.json`, 'toggle-chat', true)
+                            reject(new Error('Failed to locate or create Server Preferences\n\nPlease try chatting again...'))
+                        }
+
+                        // check if chat is disabled
+                        else if (!config.options['toggle-chat'])
+                            reject(new Error('Admin(s) have disabled chat features.\n\n Please contact your server\'s admin(s).'))
+                        else
+                            resolve(config)
+                    })
+                })
+                break // successful
+            } catch (error) {   
+                ++attempt
+                if (attempt < maxRetries) {
+                    log(`Attempt ${attempt} failed for Server Preferences. Retrying in ${delay}ms...`)
+                    await new Promise(ret => setTimeout(ret, delay))
+                } else 
+                    throw new Error(`Could not retrieve Server Preferences, please try chatting again...`)
+            }
+        }
 
-        // Retrieve User Preferences
-        const userConfig: UserConfig = await new Promise((resolve, reject) => {
-            getUserConfig(`${message.author.username}-config.json`, (config) => {
-                if (config === undefined) {
-                    openConfig(`${message.author.username}-config.json`, 'message-style', false)
-                    reject(new Error('No User Preferences is set up.\n\nCreating preferences file with \`message-style\` set as \`false\` for regular message style.\nPlease try chatting again.'))
-                    return
-                }
-
-                // check if there is a set capacity in config
-                if (typeof config.options['modify-capacity'] !== 'number')
-                    log(`Capacity is undefined, using default capacity of ${msgHist.capacity}.`)
-                else if (config.options['modify-capacity'] === msgHist.capacity)
-                    log(`Capacity matches config as ${msgHist.capacity}, no changes made.`)
-                else {
-                    log(`New Capacity found. Setting Context Capacity to ${config.options['modify-capacity']}.`)
-                    msgHist.capacity = config.options['modify-capacity']
-                }
-
-                // set stream state
-                shouldStream = config.options['message-stream'] as boolean || false
-
-                if (typeof config.options['switch-model'] !== 'string')
-                    reject(new Error(`No Model was set. Please set a model by running \`/switch-model <model of choice>\`.\n\nIf you do not have any models. Run \`/pull-model <model name>\`.`))
-
-                resolve(config)
-            })
-        })
+        // Reset attempts for User preferences
+        attempt = 0
+        let userConfig: UserConfig | undefined
+
+        while (attempt < maxRetries) {
+            try {
+                // Retrieve User Preferences
+                userConfig = await new Promise((resolve, reject) => {
+                    getUserConfig(`${message.author.username}-config.json`, (config) => {
+                        if (config === undefined) {
+                            openConfig(`${message.author.username}-config.json`, 'message-style', false)
+                            openConfig(`${message.author.username}-config.json`, 'switch-model', defaultModel)
+                            reject(new Error('No User Preferences is set up.\n\nCreating preferences file with \`message-style\` set as \`false\` for regular message style.\nPlease try chatting again.'))
+                            return
+                        }
+
+                        // check if there is a set capacity in config
+                        else if (typeof config.options['modify-capacity'] !== 'number')
+                            log(`Capacity is undefined, using default capacity of ${msgHist.capacity}.`)
+                        else if (config.options['modify-capacity'] === msgHist.capacity)
+                            log(`Capacity matches config as ${msgHist.capacity}, no changes made.`)
+                        else {
+                            log(`New Capacity found. Setting Context Capacity to ${config.options['modify-capacity']}.`)
+                            msgHist.capacity = config.options['modify-capacity']
+                        }
+
+                        // set stream state
+                        shouldStream = config.options['message-stream'] as boolean || false
+
+                        if (typeof config.options['switch-model'] !== 'string')
+                            reject(new Error(`No Model was set. Please set a model by running \`/switch-model <model of choice>\`.\n\nIf you do not have any models. Run \`/pull-model <model name>\`.`))
+
+                        resolve(config)
+                    })
+                })
+                break // successful
+        } catch (error) {
+            ++attempt
+            if (attempt < maxRetries) {
+                log(`Attempt ${attempt} failed for User Preferences. Retrying in ${delay}ms...`)
+                await new Promise(ret => setTimeout(ret, delay))
+            } else
+                throw new Error(`Could not retrieve User Preferences, please try chatting again...`)
+            }
+        }
 
         // need new check for "open/active" threads/channels here!
         let chatMessages: UserMessage[] = await new Promise((resolve) => {
@@ -106,6 +137,9 @@ export default event(Events.MessageCreate, async ({ log, msgHist, ollama, client
         // response string for ollama to put its response
         let response: string
 
+        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 messageAttachment: string[] = await getAttachmentData(message.attachments.first())
         const model: string = userConfig.options['switch-model']

diff --git a/src/keys.ts b/src/keys.ts
@@ -4,6 +4,7 @@ export const Keys = {
     clientToken: getEnvVar('CLIENT_TOKEN'),
     ipAddress: getEnvVar('OLLAMA_IP', '127.0.0.1'), // default ollama ip if none
     portAddress: getEnvVar('OLLAMA_PORT', '11434'), // default ollama port if none
+    defaultModel: getEnvVar('MODEL', 'llama3.2')
 } as const // readonly keys
 
 export default Keys
diff --git a/src/utils/events.ts b/src/utils/events.ts
@@ -36,7 +36,8 @@ export interface EventProps {
     client: Client
     log: LogMethod
     msgHist: Queue<UserMessage>
-    ollama: Ollama
+    ollama: Ollama,
+    defaultModel: String
 }
 export type EventCallback<T extends EventKeys> = (
     props: EventProps,
@@ -64,7 +65,8 @@ export function registerEvents(
     client: Client, 
     events: Event[], 
     msgHist: Queue<UserMessage>,
-    ollama: Ollama
+    ollama: Ollama,
+    defaultModel: String
 ): void {
     for (const { key, callback } of events) {
         client.on(key, (...args) => {
@@ -73,7 +75,7 @@ export function registerEvents(
 
             // Handle Errors, call callback, log errors as needed
             try {
-                callback({ client, log, msgHist, ollama }, ...args)
+                callback({ client, log, msgHist, ollama, defaultModel }, ...args)
             } catch (error) {
                 log('[Uncaught Error]', error)
             }

diff --git a/src/utils/messageNormal.ts b/src/utils/messageNormal.ts
@@ -73,7 +73,10 @@ export async function normalMessage(
             }            
         } catch(error: any) {
             console.log(`[Util: messageNormal] Error creating message: ${error.message}`)
-            sentMessage.edit(`**Response generation failed.**\n\nReason: ${error.message}`)
+            if (error.message.includes('try pulling it first'))
+                sentMessage.edit(`**Response generation failed.**\n\nReason: You do not have the ${model} downloaded. Ask an admin to pull it using the \`pull-model\` command.`)
+            else
+                sentMessage.edit(`**Response generation failed.**\n\nReason: ${error.message}`)
         }
     })