226 lines
8.3 KiB
JavaScript
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}`);
|
|
}
|
|
}
|
|
}
|