wickr-bedrock-bot
Version:
AWS Wickr's own Bedrock Bot
199 lines (198 loc) • 7.92 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.BedrockLLMClient = void 0;
const client_bedrock_runtime_1 = require("@aws-sdk/client-bedrock-runtime");
const index_js_1 = require("@modelcontextprotocol/sdk/client/index.js");
const stdio_js_1 = require("@modelcontextprotocol/sdk/client/stdio.js");
const logger_1 = __importDefault(require("./logger"));
const config_1 = require("./config");
const cachePoint = {
cachePoint: { type: 'default' },
};
class BedrockLLMClient {
bedrock;
mcpClient = null;
availableTools = [];
invocationCount = 0;
clientConfig;
constructor() {
this.bedrock = new client_bedrock_runtime_1.BedrockRuntimeClient();
this.clientConfig = this.initConfig();
this.initialize().catch((error) => {
logger_1.default.error({ error }, 'Error initializing MCP client');
});
}
initConfig() {
const config = new config_1.Configuration();
return config_1.Configuration.instance;
}
async initialize() {
const transport = new stdio_js_1.StdioClientTransport({
command: 'npx',
args: ['-y', '@modelcontextprotocol/server-aws-kb-retrieval'],
env: process.env,
});
this.mcpClient = new index_js_1.Client({ name: 'wickr-bot', version: '1.0.0' }, { capabilities: { resources: {}, tools: {} } });
await this.mcpClient.connect(transport);
const toolsResult = await this.mcpClient.listTools();
this.availableTools = toolsResult.tools.map((tool) => ({
toolSpec: {
name: tool.name,
description: tool.description,
inputSchema: { json: tool.inputSchema },
},
}));
logger_1.default.info({ toolCount: this.availableTools.length }, 'MCP client initialized with tools');
}
getInvocationCount() {
return this.invocationCount;
}
hasAvailableTools() {
return this.availableTools.length > 0;
}
async chat(options) {
const { messages, systemPrompt, modelConfig } = options;
if (modelConfig.cacheSupport?.system) {
systemPrompt.push(cachePoint);
}
const useTools = this.availableTools.length > 0 && modelConfig.supportsTools;
logger_1.default.info("Sending message to bedrock client");
return useTools
? await this.handleToolBasedConversation(modelConfig, messages, systemPrompt)
: await this.handleStandardConversation(modelConfig, messages, systemPrompt);
}
async handleToolBasedConversation(modelConfig, initialMessages, systemPrompt) {
let currentMessages = initialMessages;
const tools = modelConfig.cacheSupport?.tools
? [...this.availableTools, cachePoint]
: this.availableTools;
while (true) {
const response = await this.callModel({
modelId: modelConfig.id,
systemPrompt,
messages: currentMessages,
toolConfig: { tools },
});
this.invocationCount++;
const outputMessage = response.output.message;
currentMessages = [...currentMessages, outputMessage];
if (response.stopReason !== 'tool_use') {
return this.extractTextResponse(outputMessage);
}
await this.processToolRequests(outputMessage, currentMessages);
}
}
async processToolRequests(outputMessage, currentMessages) {
logger_1.default.debug('Model requested tool use');
for (const content of outputMessage.content) {
if (!content.toolUse)
continue;
const toolUse = content.toolUse;
logger_1.default.debug({ toolName: toolUse.name }, 'Tool requested');
try {
if (!this.mcpClient) {
throw new Error('MCP client not initialized');
}
const result = await this.mcpClient.callTool({
name: toolUse.name,
arguments: toolUse.input,
});
currentMessages.push(this.createToolResultMessage(toolUse.toolUseId, result));
}
catch (error) {
logger_1.default.error({ error, toolName: toolUse.name }, 'Tool call failed');
currentMessages.push(this.createToolErrorMessage(toolUse.toolUseId, error));
}
}
}
createToolResultMessage(toolUseId, result) {
return {
role: client_bedrock_runtime_1.ConversationRole.USER,
content: [
{
toolResult: {
toolUseId: toolUseId,
content: [{ text: JSON.stringify(result) }],
},
},
],
};
}
createToolErrorMessage(toolUseId, error) {
return {
role: client_bedrock_runtime_1.ConversationRole.USER,
content: [
{
toolResult: {
toolUseId: toolUseId,
content: [
{
text: `Error: ${error.message || String(error)}`,
},
],
status: 'error',
},
},
],
};
}
extractTextResponse(outputMessage) {
const textContents = outputMessage.content
.filter((c) => c.text)
.map((c) => c.text);
return textContents.join('\n').trim();
}
async handleStandardConversation(modelConfig, messages, systemPrompt) {
return await this.callModel({
modelId: modelConfig.id,
systemPrompt,
messages,
});
}
async callModel(options) {
const { modelId, systemPrompt, messages, toolConfig } = options;
const useInvokeModelCommands = async () => {
const input = {
contentType: "application/json",
body: JSON.stringify(`${messages}\n${systemPrompt}`),
modelId,
};
const command = new client_bedrock_runtime_1.InvokeModelCommand(input);
return this.bedrock.send(command);
};
const useConverseModelCommands = async () => {
const input = {
modelId,
messages,
system: systemPrompt,
inferenceConfig: {
maxTokens: this.clientConfig.appConfig.maxTokens,
temperature: this.clientConfig.appConfig.temperature,
},
};
if (toolConfig && this.availableTools.length > 0) {
input.toolConfig = toolConfig;
}
const command = new client_bedrock_runtime_1.ConverseCommand(input);
;
return this.bedrock.send(command);
};
const response = modelId.includes('titan') ? await useInvokeModelCommands() : await useConverseModelCommands();
if (toolConfig) {
return response;
}
const output = ("body" in response) ? JSON.parse(Buffer.from(response.body).toString('utf-8')) : response.output;
if (!output?.message?.content?.length) {
throw new Error('Invalid response from Converse API');
}
const responseText = output.message.content[0].text;
if (typeof responseText !== 'string') {
throw new Error('Invalid response text from Converse API');
}
return responseText.trim();
}
}
exports.BedrockLLMClient = BedrockLLMClient;