diff --git a/src/client.ts b/src/client.ts index 21337d9..841abe7 100644 --- a/src/client.ts +++ b/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}`, }) diff --git a/src/commands/index.ts b/src/commands/index.ts index e93a108..d81cead 100644 --- a/src/commands/index.ts +++ b/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[] \ No newline at end of file diff --git a/src/commands/pullModel.ts b/src/commands/pullModel.ts new file mode 100644 index 0000000..d35edc4 --- /dev/null +++ b/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.` + }) + } +} \ No newline at end of file diff --git a/tests/commands.test.ts b/tests/commands.test.ts index 2589048..4103395 100644 --- a/tests/commands.test.ts +++ b/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') }) }) diff --git a/tests/events.test.ts b/tests/events.test.ts index 34916a0..06414d6 100644 --- a/tests/events.test.ts +++ b/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