agsamantha/node_modules/@langchain/community/dist/stores/message/dynamodb.js
2024-10-02 15:15:21 -05:00

226 lines
8.3 KiB
JavaScript

import { DynamoDBClient, GetItemCommand, UpdateItemCommand, DeleteItemCommand, } from "@aws-sdk/client-dynamodb";
import { BaseListChatMessageHistory } from "@langchain/core/chat_history";
import { mapChatMessagesToStoredMessages, mapStoredMessagesToChatMessages, } from "@langchain/core/messages";
/**
* Class providing methods to interact with a DynamoDB table to store and
* retrieve chat messages. It extends the `BaseListChatMessageHistory`
* class.
*/
export class DynamoDBChatMessageHistory extends BaseListChatMessageHistory {
get lc_secrets() {
return {
"config.credentials.accessKeyId": "AWS_ACCESS_KEY_ID",
"config.credentials.secretAccessKey": "AWS_SECRETE_ACCESS_KEY",
"config.credentials.sessionToken": "AWS_SESSION_TOKEN",
};
}
/**
* Transforms a `StoredMessage` into a `DynamoDBSerializedChatMessage`.
* The `DynamoDBSerializedChatMessage` format is suitable for storing in DynamoDB.
*
* @param message - The `StoredMessage` to be transformed.
* @returns The transformed `DynamoDBSerializedChatMessage`.
*/
createDynamoDBSerializedChatMessage(message) {
const { type, data: { content, role, additional_kwargs }, } = message;
const isAdditionalKwargs = additional_kwargs && Object.keys(additional_kwargs).length;
const dynamoSerializedMessage = {
M: {
type: {
S: type,
},
text: {
S: content,
},
additional_kwargs: isAdditionalKwargs
? { S: JSON.stringify(additional_kwargs) }
: { S: "{}" },
},
};
if (role) {
dynamoSerializedMessage.M.role = { S: role };
}
return dynamoSerializedMessage;
}
constructor({ tableName, sessionId, partitionKey, sortKey, messageAttributeName, config, key = {}, }) {
super();
Object.defineProperty(this, "lc_namespace", {
enumerable: true,
configurable: true,
writable: true,
value: ["langchain", "stores", "message", "dynamodb"]
});
Object.defineProperty(this, "tableName", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "sessionId", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "client", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "partitionKey", {
enumerable: true,
configurable: true,
writable: true,
value: "id"
});
Object.defineProperty(this, "sortKey", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "messageAttributeName", {
enumerable: true,
configurable: true,
writable: true,
value: "messages"
});
Object.defineProperty(this, "dynamoKey", {
enumerable: true,
configurable: true,
writable: true,
value: {}
});
this.tableName = tableName;
this.sessionId = sessionId;
this.client = new DynamoDBClient(config ?? {});
this.partitionKey = partitionKey ?? this.partitionKey;
this.sortKey = sortKey;
this.messageAttributeName =
messageAttributeName ?? this.messageAttributeName;
this.dynamoKey = key;
// override dynamoKey with partition key and sort key when key not specified
if (Object.keys(this.dynamoKey).length === 0) {
this.dynamoKey[this.partitionKey] = { S: this.sessionId };
if (this.sortKey) {
this.dynamoKey[this.sortKey] = { S: this.sortKey };
}
}
}
/**
* Retrieves all messages from the DynamoDB table and returns them as an
* array of `BaseMessage` instances.
* @returns Array of stored messages
*/
async getMessages() {
try {
const params = {
TableName: this.tableName,
Key: this.dynamoKey,
};
const response = await this.client.send(new GetItemCommand(params));
const items = response.Item
? response.Item[this.messageAttributeName]?.L ?? []
: [];
const messages = items
.filter((item) => item.M !== undefined)
.map((item) => {
const data = {
role: item.M?.role?.S,
content: item.M?.text.S,
additional_kwargs: item.M?.additional_kwargs?.S
? JSON.parse(item.M?.additional_kwargs.S)
: undefined,
};
return {
type: item.M?.type.S,
data,
};
})
.filter((x) => x.type !== undefined && x.data.content !== undefined);
return mapStoredMessagesToChatMessages(messages);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
}
catch (error) {
throw new Error(`Error getting messages: ${error.message}`);
}
}
/**
* Adds a new message to the DynamoDB table.
* @param message The message to be added to the DynamoDB table.
*/
async addMessage(message) {
try {
const messages = mapChatMessagesToStoredMessages([message]);
const params = {
TableName: this.tableName,
Key: this.dynamoKey,
ExpressionAttributeNames: {
"#m": this.messageAttributeName,
},
ExpressionAttributeValues: {
":empty_list": {
L: [],
},
":m": {
L: messages.map(this.createDynamoDBSerializedChatMessage),
},
},
UpdateExpression: "SET #m = list_append(if_not_exists(#m, :empty_list), :m)",
};
await this.client.send(new UpdateItemCommand(params));
// eslint-disable-next-line @typescript-eslint/no-explicit-any
}
catch (error) {
throw new Error(`Error adding message: ${error.message}`);
}
}
/**
* Adds new messages to the DynamoDB table.
* @param messages The messages to be added to the DynamoDB table.
*/
async addMessages(messages) {
try {
const storedMessages = mapChatMessagesToStoredMessages(messages);
const dynamoMessages = storedMessages.map(this.createDynamoDBSerializedChatMessage);
const params = {
TableName: this.tableName,
Key: this.dynamoKey,
ExpressionAttributeNames: {
"#m": this.messageAttributeName,
},
ExpressionAttributeValues: {
":empty_list": {
L: [],
},
":m": {
L: dynamoMessages,
},
},
UpdateExpression: "SET #m = list_append(if_not_exists(#m, :empty_list), :m)",
};
await this.client.send(new UpdateItemCommand(params));
// eslint-disable-next-line @typescript-eslint/no-explicit-any
}
catch (error) {
throw new Error(`Error adding messages: ${error.message}`);
}
}
/**
* Deletes all messages from the DynamoDB table.
*/
async clear() {
try {
const params = {
TableName: this.tableName,
Key: this.dynamoKey,
};
await this.client.send(new DeleteItemCommand(params));
// eslint-disable-next-line @typescript-eslint/no-explicit-any
}
catch (error) {
throw new Error(`Error clearing messages: ${error.message}`);
}
}
}