Merge pull request #1697 from Jaredude/bugfix/dynamob-memory-store

DynamoDB Chat Memory fix
This commit is contained in:
Henry Heng
2024-02-09 01:36:38 +08:00
committed by GitHub
@@ -117,7 +117,10 @@ const initalizeDynamoDB = async (nodeData: INodeData, options: ICommonObject): P
memoryKey: memoryKey ?? 'chat_history',
chatHistory: dynamoDb,
sessionId,
dynamodbClient: client
dynamodbClient: client,
tableName,
partitionKey,
dynamoKey: { [partitionKey]: { S: sessionId } }
})
return memory
}
@@ -125,6 +128,9 @@ const initalizeDynamoDB = async (nodeData: INodeData, options: ICommonObject): P
interface BufferMemoryExtendedInput {
dynamodbClient: DynamoDBClient
sessionId: string
tableName: string
partitionKey: string
dynamoKey: Record<string, AttributeValue>
}
interface DynamoDBSerializedChatMessage {
@@ -142,6 +148,10 @@ interface DynamoDBSerializedChatMessage {
}
class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
private tableName = ''
private partitionKey = ''
private dynamoKey: Record<string, AttributeValue>
private messageAttributeName: string
sessionId = ''
dynamodbClient: DynamoDBClient
@@ -149,11 +159,14 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
super(fields)
this.sessionId = fields.sessionId
this.dynamodbClient = fields.dynamodbClient
this.tableName = fields.tableName
this.partitionKey = fields.partitionKey
this.dynamoKey = fields.dynamoKey
}
overrideDynamoKey(overrideSessionId = '') {
const existingDynamoKey = (this as any).dynamoKey
const partitionKey = (this as any).partitionKey
const existingDynamoKey = this.dynamoKey
const partitionKey = this.partitionKey
let newDynamoKey: Record<string, AttributeValue> = {}
@@ -209,9 +222,9 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
async getChatMessages(overrideSessionId = '', returnBaseMessages = false): Promise<IMessage[] | BaseMessage[]> {
if (!this.dynamodbClient) return []
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : (this as any).dynamoKey
const tableName = (this as any).tableName
const messageAttributeName = (this as any).messageAttributeName
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : this.dynamoKey
const tableName = this.tableName
const messageAttributeName = this.messageAttributeName
const params: GetItemCommandInput = {
TableName: tableName,
@@ -236,9 +249,9 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
async addChatMessages(msgArray: { text: string; type: MessageType }[], overrideSessionId = ''): Promise<void> {
if (!this.dynamodbClient) return
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : (this as any).dynamoKey
const tableName = (this as any).tableName
const messageAttributeName = (this as any).messageAttributeName
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : this.dynamoKey
const tableName = this.tableName
const messageAttributeName = this.messageAttributeName
const input = msgArray.find((msg) => msg.type === 'userMessage')
const output = msgArray.find((msg) => msg.type === 'apiMessage')
@@ -259,8 +272,8 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
async clearChatMessages(overrideSessionId = ''): Promise<void> {
if (!this.dynamodbClient) return
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : (this as any).dynamoKey
const tableName = (this as any).tableName
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : this.dynamoKey
const tableName = this.tableName
const params: DeleteItemCommandInput = {
TableName: tableName,