wickr-bedrock-bot
Version:
AWS Wickr's own Bedrock Bot
330 lines (329 loc) • 14.2 kB
JavaScript
"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.BedrockBot = void 0;
const wickrio_bot_framework_1 = __importStar(require("wickrio-bot-framework"));
const client_bedrock_runtime_1 = require("@aws-sdk/client-bedrock-runtime");
const zlib_1 = require("zlib");
const util_1 = require("util");
const crypto_1 = require("crypto");
const logger_1 = __importDefault(require("./logger"));
const config_1 = require("./config");
const button_helper_1 = __importDefault(require("./helpers/button-helper"));
class BedrockBot extends wickrio_bot_framework_1.default {
llmClient;
state;
startTime = Date.now();
messagesSent = 0;
botConfig;
constructor(wickr, llmClient, username) {
super(wickr, username);
this.llmClient = llmClient;
this.state = {};
this.helpText = config_1.Configuration.HELP_TEXT;
this.botConfig = this.initConfig();
this.addHandlers();
this.on('start', async () => {
this.clearStats();
this.state = await this.loadState();
setInterval(this.saveState.bind(this), config_1.Configuration.STATE_SAVE_INTERVAL_MS);
try {
this.setAvatar(__dirname + '/../../img/avatar.png');
}
catch (error) {
logger_1.default.warn({ error }, 'Error setting avatar');
}
});
}
initConfig() {
const config = new config_1.Configuration();
return config_1.Configuration.instance;
}
addHandlers() {
this.listen('chat', this.chat.bind(this), {
description: 'Message the AI assistant. This is the default behavior in 1:1 messages with the bot.',
});
this.listen('clear', this.clear.bind(this), {
description: 'Clear message context for this room',
});
this.listen('model', this.handleModelCommands.bind(this), {
description: 'Set the AI model to use for this room. Usage: /model [model-name]',
});
this.listen('stats', this.sendStats.bind(this), { hidden: true });
delete this.handlers['help'];
this.setDefaultHandler(this.chatIfDirectMessage.bind(this));
}
async _send(vgroupid, msg) {
this.messagesSent++;
await this.send(vgroupid, msg);
}
async chatIfDirectMessage(msg, args) {
try {
this.getHistory(msg.vgroupid);
this.getModelConfig(msg.vgroupid);
if (msg.msgtype !== wickrio_bot_framework_1.MESSAGE_TYPE.TEXT) {
return;
}
else if (msg.message && msg.message.startsWith('/')) {
await this.send(msg.vgroupid, this.helpText?.trim() || '');
return;
}
else if (msg.is_room) {
if (args)
this.addMessageToHistory(msg, args.join(' '));
return;
}
else {
await this.chat(msg, args);
}
}
catch (error) {
logger_1.default.error({ error }, 'Error while processing direct message to bot');
}
}
async handleModelCommands(msg, args) {
logger_1.default.debug({ args }, "Args in the model command");
try {
this.getHistory(msg.vgroupid);
this.getModelConfig(msg.vgroupid);
const modelsList = Object.keys(this.botConfig.modelsConfig.models);
if (msg.is_room) {
if (msg.message === '/model') {
await this.send(msg.vgroupid, `I am currently running model name \`${this.state[msg.vgroupid].modelName}\` If you want to change to a different model, type /model <new model name>. Here is a list of available models: ${Object.keys(this.botConfig.modelsConfig.models).join(', ')}`);
}
else if (msg.message && msg.message.split(' ').length === 2) {
const parsedMessage = msg.message.replace(/(?:\\(.))/g, '$1');
const requestedModelName = parsedMessage.split(' ')[1];
logger_1.default.debug({ requestedModelName }, 'Requested Model Name');
this.setModel(msg, [requestedModelName]);
}
}
else {
if (msg.message === '/model') {
const buttons = button_helper_1.default.makeYesNoButton({ yesButtonAction: '/model list', noButtonAction: undefined });
await this.send(msg.vgroupid, `I am currently running ${this.state[msg.vgroupid].modelName} model. Do you want to change this?`, {
meta: {
...buttons
}
});
}
else if (msg.message === '/model list') {
const rows = modelsList.map((modelName) => {
return {
firstcolvalue: modelName,
response: `/model ${modelName}`
};
});
const meta = {
table: {
name: 'List of Models',
firstcolname: 'Model Name',
actioncolname: 'Select Model',
rows,
},
textcut: [
{
startindex: 0,
endindex: 15,
},
],
};
await this.send(msg.vgroupid, 'Select the list of models from the table below:\n', {
meta
});
}
else if (msg.message && msg.message.split(' ').length === 2) {
const parsedMessage = msg.message.replace(/(?:\\(.))/g, '$1');
const requestedModelName = parsedMessage.split(' ')[1];
logger_1.default.debug({ requestedModelName }, 'Requested Model Name');
this.setModel(msg, [requestedModelName]);
}
}
}
catch (error) {
logger_1.default.error({ error }, 'Error while processing model command');
}
}
getHistory(vgroupid) {
if (!this.state[vgroupid]) {
this.state[vgroupid] = { history: [] };
}
return this.state[vgroupid].history;
}
addMessageToHistory(msg, content, role = client_bedrock_runtime_1.ConversationRole.USER) {
const { vgroupid } = msg;
const history = this.getHistory(vgroupid);
history.push({
role,
content,
metadata: {
timestamp: msg.time,
sender: msg.sender,
},
});
if (history.length > this.botConfig.appConfig.config.maxHistory) {
history.splice(0, 1);
}
return history;
}
getModelConfig(vgroupid) {
const currentModelName = this.state[vgroupid].modelName;
if (currentModelName) {
return this.botConfig.modelsConfig.models[currentModelName];
}
else {
const modelName = this.botConfig.modelsConfig.defaultModel;
this.state[vgroupid].modelName = modelName;
return this.botConfig.modelsConfig.models[modelName];
}
}
convertHistoryToPromptMessages(history) {
return history.map((message) => {
const { metadata, content } = message;
const text = `${config_1.PROMPT_METADATA_START_TAG}${JSON.stringify(metadata)}${config_1.PROMPT_METADATA_END_TAG}\n${content}`;
return {
role: message.role,
content: [{ text }],
};
});
}
async chat(msg, args) {
const userMessage = args ? args.join(' ') : '';
const { vgroupid } = msg;
const history = this.addMessageToHistory(msg, userMessage);
const messages = this.convertHistoryToPromptMessages(history);
const modelConfig = this.getModelConfig(vgroupid);
try {
const systemPrompt = this.getSystemContext(msg);
let finalResponse = await this.llmClient.chat({
messages,
systemPrompt,
modelConfig,
});
if (finalResponse.startsWith(config_1.PROMPT_METADATA_START_TAG)) {
logger_1.default.warn('LLM returned metadata in response');
finalResponse = finalResponse.split(config_1.PROMPT_METADATA_END_TAG + '\n')[1];
}
this.addMessageToHistory(msg, finalResponse, client_bedrock_runtime_1.ConversationRole.ASSISTANT);
await this._send(vgroupid, finalResponse);
}
catch (error) {
const requestId = (0, crypto_1.randomUUID)();
let msg = 'Sorry, I encountered an error completing your request.';
logger_1.default.error({ error, requestId }, 'Error calling model');
if (error instanceof client_bedrock_runtime_1.ThrottlingException) {
msg =
'Sorry, my request was rate limited when trying to generate a response for you. Please try again later.';
}
await this._send(vgroupid, msg + ` Request ID: ${requestId}`);
}
}
setModel(msg, args) {
const { vgroupid } = msg;
const modelName = args[0];
if (!modelName) {
const config = this.getModelConfig(vgroupid);
this._send(vgroupid, `Model successfully changed! I am now running on ${config.id}`);
return;
}
if (!(Object.keys(this.botConfig.modelsConfig.models).includes(modelName))) {
const availableModels = Object.keys(this.botConfig.modelsConfig.models).join(', ');
this._send(vgroupid, ` This is not a valid model name. These are the models available to you: ${availableModels}`);
return;
}
if (!this.state[vgroupid]) {
this.state[vgroupid] = { history: [] };
}
this.state[vgroupid].modelName = modelName;
const modelConfig = this.botConfig.modelsConfig.models[modelName];
this._send(vgroupid, `Model set to ${modelName} (${modelConfig.id})`);
}
clear(msg, args) {
const { vgroupid } = msg;
const roomState = this.state[vgroupid];
if (roomState) {
roomState.history = [];
this._send(vgroupid, 'The context for this room has been reset');
}
}
sendStats(msg) {
if (msg.is_room)
return;
const stats = `- State size: ${JSON.stringify(this.state).length}
- LLM invocations: ${this.llmClient.getInvocationCount()}
- Messages sent: ${this.messagesSent}
- Uptime (seconds): ${(Date.now() - this.startTime) / 1000}`;
this._send(msg.vgroupid, stats);
}
getSystemContext(msg) {
const context = msg.is_room ? config_1.Configuration.getRoomContext() : config_1.Configuration.getDmContext();
const timestamp = new Date().toISOString().replace('T', ' ').slice(0, -1);
return [{ text: context + `\n\nThe current time is ${timestamp} UTC` }];
}
async loadState() {
try {
const data = await this.brain.get(config_1.Configuration.STATE_PATH);
if (data) {
const compressedBuffer = Buffer.from(data, 'base64');
const jsonBuffer = await (0, util_1.promisify)(zlib_1.gunzip)(compressedBuffer);
return JSON.parse(jsonBuffer.toString('utf-8'));
}
}
catch (error) {
logger_1.default.error({ error }, 'Error loading saved state');
}
return {};
}
async saveState() {
const jsonString = JSON.stringify(this.state);
logger_1.default.debug({ stateSize: jsonString.length }, 'Preparing state for database');
try {
const compressedBuffer = await (0, util_1.promisify)(zlib_1.gzip)(Buffer.from(jsonString, 'utf-8'));
const base64String = compressedBuffer.toString('base64');
logger_1.default.debug({
compressedSize: base64String.length,
compressionRatio: Math.round((base64String.length / jsonString.length) * 100),
}, 'Saving compressed state to database');
this.brain.set(config_1.Configuration.STATE_PATH, base64String);
}
catch (error) {
logger_1.default.error({ error }, 'Error updating state');
}
}
}
exports.BedrockBot = BedrockBot;