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

…-store DynamoDB Chat Memory fix
georggr · Feb 8, 2024 · 1f2e58c · 1f2e58c
2 parents 64018b7 + dd89af8
commit 1f2e58c
Showing 1 changed file with 24 additions and 11 deletions.
diff --git a/packages/components/nodes/memory/DynamoDb/DynamoDb.ts b/packages/components/nodes/memory/DynamoDb/DynamoDb.ts
@@ -117,14 +117,20 @@ 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
 }
 
 interface BufferMemoryExtendedInput {
     dynamodbClient: DynamoDBClient
     sessionId: string
+    tableName: string
+    partitionKey: string
+    dynamoKey: Record<string, AttributeValue>
 }
 
 interface DynamoDBSerializedChatMessage {
@@ -142,18 +148,25 @@ interface DynamoDBSerializedChatMessage {
 }
 
 class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
+    private tableName = ''
+    private partitionKey = ''
+    private dynamoKey: Record<string, AttributeValue>
+    private messageAttributeName: string
     sessionId = ''
     dynamodbClient: DynamoDBClient
 
     constructor(fields: BufferMemoryInput & BufferMemoryExtendedInput) {
         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,