From d67106c03e2356ff3426c506f692442a99e70691 Mon Sep 17 00:00:00 2001 From: Kevin Dang <77701718+kevinthedang@users.noreply.github.com> Date: Sun, 21 Apr 2024 14:40:30 -0700 Subject: [PATCH] Chat Stream Integration (#52) * rm: axios dependency * add: stream parsing for normal style * fix: empty string problem * add: stream for embedded prompts * update: version increment --- docker-compose.yml | 2 +- package-lock.json | 91 ------------------------------------- package.json | 3 +- src/events/messageCreate.ts | 13 ++++-- src/utils/events.ts | 12 +++++ src/utils/index.ts | 3 +- src/utils/messageEmbed.ts | 63 +++++++++++++++---------- src/utils/messageNormal.ts | 59 ++++++++++++++---------- src/utils/streamHandler.ts | 40 ++++++++++++++++ src/utils/streamParse.ts | 27 ----------- 10 files changed, 140 insertions(+), 173 deletions(-) create mode 100644 src/utils/streamHandler.ts delete mode 100644 src/utils/streamParse.ts diff --git a/docker-compose.yml b/docker-compose.yml index 68f707c..752cd2a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -8,7 +8,7 @@ services: build: ./ # find docker file in designated path container_name: discord restart: always # rebuild container always - image: discord/bot:0.4.0 + image: discord/bot:0.4.2 environment: CLIENT_TOKEN: ${CLIENT_TOKEN} GUILD_ID: ${GUILD_ID} diff --git a/package-lock.json b/package-lock.json index 29e1f37..d3744f4 100644 --- a/package-lock.json +++ b/package-lock.json @@ -9,7 +9,6 @@ "version": "0.4.0", "license": "ISC", "dependencies": { - "axios": "^1.6.2", "discord.js": "^14.14.1", "dotenv": "^16.3.1", "ollama": "^0.5.0" @@ -661,21 +660,6 @@ "integrity": "sha512-58S9QDqG0Xx27YwPSt9fJxivjYl432YCwfDMfZ+71RAqUrZef7LrKQZ3LHLOwCS4FLNBplP533Zx895SeOCHvA==", "dev": true }, - "node_modules/asynckit": { - "version": "0.4.0", - "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz", - "integrity": "sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==" - }, - "node_modules/axios": { - "version": "1.6.8", - "resolved": "https://registry.npmjs.org/axios/-/axios-1.6.8.tgz", - "integrity": "sha512-v/ZHtJDU39mDpyBoFVkETcd/uNdxrWRrg3bKpOKzXFA6Bvqopts6ALSMU3y6ijYxbw2B+wPrIv46egTzJXCLGQ==", - "dependencies": { - "follow-redirects": "^1.15.6", - "form-data": "^4.0.0", - "proxy-from-env": "^1.1.0" - } - }, "node_modules/balanced-match": { "version": "1.0.2", "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz", @@ -740,17 +724,6 @@ "fsevents": "~2.3.2" } }, - "node_modules/combined-stream": { - "version": "1.0.8", - "resolved": "https://registry.npmjs.org/combined-stream/-/combined-stream-1.0.8.tgz", - "integrity": "sha512-FQN4MRfuJeHf7cBbBMJFXhKSDq+2kAArBlmRBvcvFE5BB1HZKXtSFASDhdlz9zOYwxh8lDdnvmMOe/+5cdoEdg==", - "dependencies": { - "delayed-stream": "~1.0.0" - }, - "engines": { - "node": ">= 0.8" - } - }, "node_modules/concat-map": { "version": "0.0.1", "resolved": "https://registry.npmjs.org/concat-map/-/concat-map-0.0.1.tgz", @@ -780,14 +753,6 @@ } } }, - "node_modules/delayed-stream": { - "version": "1.0.0", - "resolved": "https://registry.npmjs.org/delayed-stream/-/delayed-stream-1.0.0.tgz", - "integrity": "sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==", - "engines": { - "node": ">=0.4.0" - } - }, "node_modules/diff": { "version": "4.0.2", "resolved": "https://registry.npmjs.org/diff/-/diff-4.0.2.tgz", @@ -892,38 +857,6 @@ "node": ">=8" } }, - "node_modules/follow-redirects": { - "version": "1.15.6", - "resolved": "https://registry.npmjs.org/follow-redirects/-/follow-redirects-1.15.6.tgz", - "integrity": "sha512-wWN62YITEaOpSK584EZXJafH1AGpO8RVgElfkuXbTOrPX4fIfOyEpW/CsiNd8JdYrAoOvafRTOEnvsO++qCqFA==", - "funding": [ - { - "type": "individual", - "url": "https://github.com/sponsors/RubenVerborgh" - } - ], - "engines": { - "node": ">=4.0" - }, - "peerDependenciesMeta": { - "debug": { - "optional": true - } - } - }, - "node_modules/form-data": { - "version": "4.0.0", - "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.0.tgz", - "integrity": "sha512-ETEklSGi5t0QMZuiXoA/Q6vcnxcLQP5vdugSpuAyi6SVGi2clPPp+xgEhuMaHC+zGgn31Kd235W35f7Hykkaww==", - "dependencies": { - "asynckit": "^0.4.0", - "combined-stream": "^1.0.8", - "mime-types": "^2.1.12" - }, - "engines": { - "node": ">= 6" - } - }, "node_modules/fsevents": { "version": "2.3.3", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz", @@ -1052,25 +985,6 @@ "integrity": "sha512-s8UhlNe7vPKomQhC1qFelMokr/Sc3AgNbso3n74mVPA5LTZwkB9NlXf4XPamLxJE8h0gh73rM94xvwRT2CVInw==", "dev": true }, - "node_modules/mime-db": { - "version": "1.52.0", - "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.52.0.tgz", - "integrity": "sha512-sPU4uV7dYlvtWJxwwxHD0PuihVNiE7TyAbQ5SWxDCB9mUYvOgroQOwYQQOKPJ8CIbE+1ETVlOoK1UC2nU3gYvg==", - "engines": { - "node": ">= 0.6" - } - }, - "node_modules/mime-types": { - "version": "2.1.35", - "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-2.1.35.tgz", - "integrity": "sha512-ZDY+bPm5zTTF+YpCrAU9nK0UgICYPT0QtT1NZWFv4s++TNkcgVaT0g6+4R2uI4MjQjzysHB1zxuWL50hzaeXiw==", - "dependencies": { - "mime-db": "1.52.0" - }, - "engines": { - "node": ">= 0.6" - } - }, "node_modules/minimatch": { "version": "3.1.2", "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.2.tgz", @@ -1161,11 +1075,6 @@ "url": "https://github.com/sponsors/jonschlinkert" } }, - "node_modules/proxy-from-env": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/proxy-from-env/-/proxy-from-env-1.1.0.tgz", - "integrity": "sha512-D+zkORCbA9f1tdWRK0RaCR3GPv50cMxcrz4X8k5LTSUD1Dkw47mKJEZQNunItRTkWwgtaUSo1RVFRIG9ZXiFYg==" - }, "node_modules/pstree.remy": { "version": "1.1.8", "resolved": "https://registry.npmjs.org/pstree.remy/-/pstree.remy-1.1.8.tgz", diff --git a/package.json b/package.json index df08fe7..28a432a 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "discord-ollama", - "version": "0.4.0", + "version": "0.4.2", "description": "Ollama Integration into discord", "main": "build/index.js", "exports": "./build/index.js", @@ -24,7 +24,6 @@ "author": "Kevin Dang", "license": "ISC", "dependencies": { - "axios": "^1.6.2", "discord.js": "^14.14.1", "dotenv": "^16.3.1", "ollama": "^0.5.0" diff --git a/src/events/messageCreate.ts b/src/events/messageCreate.ts index 54234ba..3ed6877 100644 --- a/src/events/messageCreate.ts +++ b/src/events/messageCreate.ts @@ -19,6 +19,8 @@ export default event(Events.MessageCreate, async ({ log, msgHist, tokens, ollama // Only respond if message mentions the bot if (!message.mentions.has(tokens.clientUid)) return + let shouldStream = false + // Try to query and send embed try { const config: Configuration = await new Promise((resolve, reject) => { @@ -45,11 +47,14 @@ export default event(Events.MessageCreate, async ({ log, msgHist, tokens, ollama msgHist.capacity = config.options['modify-capacity'] } + // set stream state + shouldStream = config.options['message-stream'] as boolean || false + resolve(config) }) }) - let response: ChatResponse + let response: string // check if we can push, if not, remove oldest while (msgHist.size() >= msgHist.capacity) msgHist.dequeue() @@ -62,9 +67,9 @@ export default event(Events.MessageCreate, async ({ log, msgHist, tokens, ollama // undefined or false, use normal, otherwise use embed if (config.options['message-style']) - response = await embedMessage(message, ollama, tokens, msgHist) + response = await embedMessage(message, ollama, tokens, msgHist, shouldStream) else - response = await normalMessage(message, ollama, tokens, msgHist) + response = await normalMessage(message, ollama, tokens, msgHist, shouldStream) // If something bad happened, remove user query and stop if (response == undefined) { msgHist.pop(); return } @@ -75,7 +80,7 @@ export default event(Events.MessageCreate, async ({ log, msgHist, tokens, ollama // successful query, save it in context history msgHist.enqueue({ role: 'assistant', - content: response.message.content + content: response }) } catch (error: any) { msgHist.pop() // remove message because of failure diff --git a/src/utils/events.ts b/src/utils/events.ts index 13114e5..bde6413 100644 --- a/src/utils/events.ts +++ b/src/utils/events.ts @@ -20,6 +20,18 @@ export type Tokens = { clientUid: string } +/** + * Parameters to run the chat query + * @param model the model to run + * @param ollama ollama api client + * @param msgHist message history + */ +export type ChatParams = { + model: string, + ollama: Ollama, + msgHist: UserMessage[] +} + /** * Format for the messages to be stored when communicating when the bot * @param role either assistant, user, or system diff --git a/src/utils/index.ts b/src/utils/index.ts index 178f4b4..b5ff26f 100644 --- a/src/utils/index.ts +++ b/src/utils/index.ts @@ -3,4 +3,5 @@ export * from './env.js' export * from './events.js' export * from './messageEmbed.js' export * from './messageNormal.js' -export * from './commands.js' \ No newline at end of file +export * from './commands.js' +export * from './streamHandler.js' \ No newline at end of file diff --git a/src/utils/messageEmbed.ts b/src/utils/messageEmbed.ts index 4ed9f19..69549c1 100644 --- a/src/utils/messageEmbed.ts +++ b/src/utils/messageEmbed.ts @@ -1,6 +1,6 @@ import { EmbedBuilder, Message } from 'discord.js' import { ChatResponse, Ollama } from 'ollama' -import { UserMessage } from './events.js' +import { ChatParams, UserMessage, streamResponse, blockResponse } from './index.js' import { Queue } from '../queues/queue.js' /** @@ -16,10 +16,12 @@ export async function embedMessage( channel: string, model: string }, - msgHist: Queue -) { + msgHist: Queue, + stream: boolean +): Promise { // bot response - let response: ChatResponse + let response: ChatResponse | AsyncGenerator + let result: string = '' // initial message to client const botMessage = new EmbedBuilder() @@ -30,28 +32,43 @@ export async function embedMessage( // send the message const sentMessage = await message.channel.send({ embeds: [botMessage] }) - try { - // Attempt to query model for message - response = await ollama.chat({ - model: tokens.model, - messages: msgHist.getItems(), - options: { - num_thread: 8, // remove if optimization needed further - mirostat: 1, - mirostat_tau: 2.0, - top_k: 70 - }, - stream: false - }) + // create params + const params: ChatParams = { + model: tokens.model, + ollama: ollama, + msgHist: msgHist.getItems() + } - // dummy message to let user know that query is underway - const newEmbed = new EmbedBuilder() + try { + // check if embed needs to stream + if (stream) { + response = await streamResponse(params) + + for await (const portion of response) { + result += portion.message.content + + // new embed per token... + const newEmbed = new EmbedBuilder() + .setTitle(`Responding to ${message.author.tag}`) + .setDescription(result || 'No Content Yet...') + .setColor('#00FF00') + + // edit the message + sentMessage.edit({ embeds: [newEmbed] }) + } + } else { + response = await blockResponse(params) + result = response.message.content + + // only need to create 1 embed again + const newEmbed = new EmbedBuilder() .setTitle(`Responding to ${message.author.tag}`) - .setDescription(response.message.content || 'No Content to Provide...') + .setDescription(result || 'No Content to Provide...') .setColor('#00FF00') - // edit the message - sentMessage.edit({ embeds: [newEmbed] }) + // edit the message + sentMessage.edit({ embeds: [newEmbed] }) + } } catch(error: any) { console.log(`[Util: messageEmbed] Error creating message: ${error.message}`) const errorEmbed = new EmbedBuilder() @@ -64,5 +81,5 @@ export async function embedMessage( } // Hope there is a response! undefined otherwie - return response!! + return result } \ No newline at end of file diff --git a/src/utils/messageNormal.ts b/src/utils/messageNormal.ts index 29bbdfe..3c2f55f 100644 --- a/src/utils/messageNormal.ts +++ b/src/utils/messageNormal.ts @@ -1,6 +1,6 @@ import { Message } from 'discord.js' import { ChatResponse, Ollama } from 'ollama' -import { UserMessage } from './events.js' +import { ChatParams, UserMessage, streamResponse, blockResponse } from './index.js' import { Queue } from '../queues/queue.js' /** @@ -16,38 +16,49 @@ export async function normalMessage( channel: string, model: string }, - msgHist: Queue -) { + msgHist: Queue, + stream: boolean +): Promise { // bot's respnse - let response: ChatResponse + let response: ChatResponse | AsyncGenerator + let result: string = '' await message.channel.send('Generating Response . . .').then(async sentMessage => { try { - // Attempt to query model for message - response = await ollama.chat({ + const params: ChatParams = { model: tokens.model, - messages: msgHist.getItems(), - options: { - num_thread: 8, // remove if optimization needed further - mirostat: 1, - mirostat_tau: 2.0, - top_k: 70 - }, - stream: false - }) + ollama: ollama, + msgHist: msgHist.getItems() + } - // check if message length > discord max for normal messages - if (response.message.content.length > 2000) { - sentMessage.edit(response.message.content.slice(0, 2000)) - message.channel.send(response.message.content.slice(2000)) - } else // edit the 'generic' response to new message - sentMessage.edit(response.message.content) + // run query based on stream preference, true = stream, false = block + if (stream) { + response = await streamResponse(params) + for await (const portion of response) { + // append token to message + result += portion.message.content + + // resent current output, THIS WILL BE SLOW due to discord limits! + sentMessage.edit(result || 'No Content Yet...') + } + } + else { + response = await blockResponse(params) + result = response.message.content + + // check if message length > discord max for normal messages + if (result.length > 2000) { + sentMessage.edit(result.slice(0, 2000)) + message.channel.send(result.slice(2000)) + } else // edit the 'generic' response to new message + sentMessage.edit(result) + } } catch(error: any) { console.log(`[Util: messageNormal] Error creating message: ${error.message}`) sentMessage.edit(`**Response generation failed.**\n\nReason: ${error.message}`) } }) - // Hope there is a response, force client to believe - return response!! -} \ No newline at end of file + // return the string representation of response + return result +} diff --git a/src/utils/streamHandler.ts b/src/utils/streamHandler.ts new file mode 100644 index 0000000..c205295 --- /dev/null +++ b/src/utils/streamHandler.ts @@ -0,0 +1,40 @@ +import { ChatResponse } from "ollama" +import { ChatParams } from "./index.js" + +/** + * Method to query the Ollama client for async generation + * @param params + * @returns Asyn + */ +export async function streamResponse(params: ChatParams): Promise> { + return await params.ollama.chat({ + model: params.model, + messages: params.msgHist, + options: { + num_thread: 8, // remove if optimization needed further + mirostat: 1, + mirostat_tau: 2.0, + top_k: 70 + }, + stream: true + }) +} + +/** + * Method to query the Ollama client for a block response + * @param params parameters to query the client + * @returns ChatResponse generated by the Ollama client + */ +export async function blockResponse(params: ChatParams): Promise { + return await params.ollama.chat({ + model: params.model, + messages: params.msgHist, + options: { + num_thread: 8, // remove if optimization needed further + mirostat: 1, + mirostat_tau: 2.0, + top_k: 70 + }, + stream: false + }) +} \ No newline at end of file diff --git a/src/utils/streamParse.ts b/src/utils/streamParse.ts deleted file mode 100644 index 48e27c6..0000000 --- a/src/utils/streamParse.ts +++ /dev/null @@ -1,27 +0,0 @@ -import { AxiosResponse } from 'axios' - -/** - * When running a /api/chat stream, the output needs to be parsed into an array of objects - * This method is used for development purposes and testing - * - * This will not work as intended with the inclusion of ollama-js, needs to be modified to work with it - * - * @param stream Axios response to from Ollama - */ -export function parseStream(stream: AxiosResponse) { - // split string by newline - const keywordObjects: string[] = stream.data.trim().split('\n') - - // parse string and load them into objects - const keywordsArray: { - model: string, - created_at: string, - message: { - role: string, - content: string - }, - done: boolean - }[] = keywordObjects.map((keywordString) => JSON.parse(keywordString)) - - return keywordsArray -} \ No newline at end of file