Pull Model Command (#125)

* Add: Pull Model Command

* Fix: Missing ollama mock for PullModel

src/client.ts

@@ -17,7 +17,7 @@ const client = new Client({
 });
 
 // initialize connection to ollama container
-const ollama = new Ollama({
+export const ollama = new Ollama({
     host: `http://${Keys.ipAddress}:${Keys.portAddress}`,
 })
 

src/commands/index.ts

@@ -7,6 +7,7 @@ import { Shutoff } from './shutoff.js'
 import { Capacity } from './capacity.js'
 import { PrivateThreadCreate } from './threadPrivateCreate.js'
 import { ClearUserChannelHistory } from './cleanUserChannelHistory.js'
+import { PullModel } from './pullModel.js'
 
 export default [
     ThreadCreate,
@@ -16,5 +17,6 @@ export default [
     Disable,
     Shutoff,
     Capacity,
-    ClearUserChannelHistory
+    ClearUserChannelHistory,
+    PullModel
 ] as SlashCommand[]

src/commands/pullModel.ts

@@ -0,0 +1,46 @@
+import { ApplicationCommandOptionType, ChannelType, Client, CommandInteraction } from "discord.js";
+import { SlashCommand } from "../utils/commands.js";
+import { ollama } from "../client.js";
+
+export const PullModel: SlashCommand = {
+    name: 'pull-model',
+    description: 'pulls a model from the ollama model library',
+
+    // set available user options to pass to the command
+    options: [
+        {
+            name: 'model-to-pull',
+            description: 'the name of the model to pull',
+            type: ApplicationCommandOptionType.String,
+            required: true
+        }
+    ],
+
+    // Pull for model from Ollama library
+    run: async (client: Client, interaction: CommandInteraction) => {
+        // defer reply to avoid timeout
+        await interaction.deferReply()
+
+        // fetch channel and message
+        const channel = await client.channels.fetch(interaction.channelId)
+        if (!channel || channel.type !== (ChannelType.PrivateThread && ChannelType.PublicThread && ChannelType.GuildText)) return
+
+        try {
+            // call ollama to pull desired model
+            await ollama.pull({
+                model: interaction.options.get('model-to-pull')!!.value as string
+            })
+        } catch (error) {
+            // could not resolve pull or model unfound
+            interaction.editReply({
+                content: `Could not pull/locate the **${interaction.options.get('model-to-pull')!!.value}** model within the [Ollama Model Library](https://ollama.com/library).\n\nPlease check the model library and try again.`
+            })
+            return
+        }
+
+        // successful pull
+        interaction.editReply({
+            content: `Successfully added **${interaction.options.get('model-to-pull')!!.value}** into your local model library.`
+        })
+    }
+}

tests/commands.test.ts

@@ -22,7 +22,7 @@ describe('Commands Existence', () => {
     // test specific commands in the object
     it('references specific commands', () => {
         const commandsString = commands.map(e => e.name).join(', ')
-        expect(commandsString).toBe('thread, private-thread, message-style, message-stream, toggle-chat, shutoff, modify-capacity, clear-user-channel-history')
+        expect(commandsString).toBe('thread, private-thread, message-style, message-stream, toggle-chat, shutoff, modify-capacity, clear-user-channel-history, pull-model')
     })
 })
 

tests/events.test.ts

@@ -1,13 +1,21 @@
-import { describe, expect, it } from 'vitest'
+import { describe, expect, it, vi } from 'vitest'
 import events from '../src/events/index.js'
 
+/**
+ * Mocking ollama found in client.ts because pullModel.ts
+ * relies on the existence on ollama. To prevent the mock,
+ * we will have to pass through ollama to the commands somehow.
+ */
+vi.mock('../src/client.js', () => ({
+    ollama: {
+        pull: vi.fn() // Mock the pull method found with ollama
+    }
+}))
+
 /**
  * Events test suite, tests the events object
  * Each event is to be tested elsewhere, this file
  * is to ensure that the events object is defined.
- * 
- * @param name name of the test suite
- * @param fn function holding tests to run
  */
 describe('Events Existence', () => {
     // test definition of events object