@bestdefense/bd-agent
Version:
An AI-powered coding assistant CLI that connects to AWS Bedrock
165 lines • 7.55 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.BedrockClient = void 0;
const client_bedrock_runtime_1 = require("@aws-sdk/client-bedrock-runtime");
const config_1 = require("../utils/config");
const chalk_1 = __importDefault(require("chalk"));
const error_handler_1 = require("../utils/error-handler");
class BedrockClient {
client;
modelId;
retry;
constructor() {
this.retry = (0, error_handler_1.createRetryHandler)(3, 1000);
const clientConfig = {
region: config_1.config.AWS_REGION || 'us-east-1'
};
// Only add credentials if they are properly defined
if (config_1.config.AWS_ACCESS_KEY_ID && config_1.config.AWS_SECRET_ACCESS_KEY) {
clientConfig.credentials = {
accessKeyId: config_1.config.AWS_ACCESS_KEY_ID,
secretAccessKey: config_1.config.AWS_SECRET_ACCESS_KEY
};
// Only add session token if it exists and is not empty
if (config_1.config.AWS_SESSION_TOKEN && config_1.config.AWS_SESSION_TOKEN.trim()) {
clientConfig.credentials.sessionToken = config_1.config.AWS_SESSION_TOKEN.trim();
}
}
try {
this.client = new client_bedrock_runtime_1.BedrockRuntimeClient(clientConfig);
}
catch (error) {
console.error('Failed to initialize Bedrock client:', error);
throw new Error('Invalid AWS credentials. Please check your configuration.');
}
// Support both model IDs and inference profile ARNs
this.modelId = config_1.config.BEDROCK_MODEL_ID || 'anthropic.claude-3-5-sonnet-20241022-v2:0';
// Log if using an inference profile
if (this.modelId.startsWith('arn:') || this.modelId.includes('.anthropic.')) {
console.log(chalk_1.default.gray(`Using inference profile: ${this.modelId}`));
}
}
async sendMessage(messages, tools, systemPrompt) {
try {
const formattedMessages = messages.map(msg => ({
role: msg.role,
content: typeof msg.content === 'string'
? [{ text: msg.content }]
: msg.content
}));
const input = {
modelId: this.modelId,
messages: formattedMessages,
system: systemPrompt ? [{ text: systemPrompt }] : undefined,
toolConfig: tools ? {
tools: tools.map(tool => ({
toolSpec: {
name: tool.name,
description: tool.description,
inputSchema: {
json: tool.input_schema
}
}
}))
} : undefined,
inferenceConfig: {
maxTokens: 4096,
temperature: 0.0
}
};
const command = new client_bedrock_runtime_1.ConverseCommand(input);
return await this.retry(() => this.client.send(command), error_handler_1.handleBedrockError);
}
catch (error) {
throw (0, error_handler_1.handleBedrockError)(error);
}
}
async *streamMessage(messages, tools, systemPrompt) {
try {
const formattedMessages = messages.map(msg => ({
role: msg.role,
content: typeof msg.content === 'string'
? [{ text: msg.content }]
: msg.content
}));
const input = {
modelId: this.modelId,
messages: formattedMessages,
system: systemPrompt ? [{ text: systemPrompt }] : undefined,
toolConfig: tools ? {
tools: tools.map(tool => ({
toolSpec: {
name: tool.name,
description: tool.description,
inputSchema: {
json: tool.input_schema
}
}
}))
} : undefined,
inferenceConfig: {
maxTokens: 4096,
temperature: 0.0
}
};
const command = new client_bedrock_runtime_1.ConverseStreamCommand(input);
const response = await this.retry(() => this.client.send(command), error_handler_1.handleBedrockError);
if (response.stream) {
for await (const event of response.stream) {
if (event.contentBlockDelta) {
const delta = event.contentBlockDelta.delta;
if (delta && 'text' in delta && delta.text) {
yield delta.text;
}
}
else if (event.contentBlockStart) {
const start = event.contentBlockStart.start;
if (start && 'toolUse' in start && start.toolUse) {
// For tool use, we need to accumulate the entire tool call
const toolUseStart = start.toolUse;
let inputJson = '';
// Continue reading until we get the complete tool input
for await (const nextEvent of response.stream) {
if (nextEvent.contentBlockDelta?.delta && 'toolUse' in nextEvent.contentBlockDelta.delta) {
const toolDelta = nextEvent.contentBlockDelta.delta.toolUse;
if (toolDelta?.input) {
inputJson += toolDelta.input;
}
}
else if (nextEvent.contentBlockStop) {
// Tool input is complete, create the full ToolUseBlock
let parsedInput;
try {
parsedInput = JSON.parse(inputJson);
}
catch (e) {
parsedInput = inputJson;
}
const completeToolUse = {
toolUseId: toolUseStart.toolUseId,
name: toolUseStart.name,
input: parsedInput
};
yield completeToolUse;
break;
}
}
}
}
else if (event.messageStop) {
// Stream has ended normally
return;
}
}
}
}
catch (error) {
throw (0, error_handler_1.handleBedrockError)(error);
}
}
}
exports.BedrockClient = BedrockClient;
//# sourceMappingURL=bedrock-client.js.map