Skip to content

Commit 1f2e58c

Browse files
authored
Merge pull request #1697 from Jaredude/bugfix/dynamob-memory-store
DynamoDB Chat Memory fix
2 parents 64018b7 + dd89af8 commit 1f2e58c

File tree

1 file changed

+24
-11
lines changed

1 file changed

+24
-11
lines changed

packages/components/nodes/memory/DynamoDb/DynamoDb.ts

+24-11
Original file line numberDiff line numberDiff line change
@@ -117,14 +117,20 @@ const initalizeDynamoDB = async (nodeData: INodeData, options: ICommonObject): P
117117
memoryKey: memoryKey ?? 'chat_history',
118118
chatHistory: dynamoDb,
119119
sessionId,
120-
dynamodbClient: client
120+
dynamodbClient: client,
121+
tableName,
122+
partitionKey,
123+
dynamoKey: { [partitionKey]: { S: sessionId } }
121124
})
122125
return memory
123126
}
124127

125128
interface BufferMemoryExtendedInput {
126129
dynamodbClient: DynamoDBClient
127130
sessionId: string
131+
tableName: string
132+
partitionKey: string
133+
dynamoKey: Record<string, AttributeValue>
128134
}
129135

130136
interface DynamoDBSerializedChatMessage {
@@ -142,18 +148,25 @@ interface DynamoDBSerializedChatMessage {
142148
}
143149

144150
class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
151+
private tableName = ''
152+
private partitionKey = ''
153+
private dynamoKey: Record<string, AttributeValue>
154+
private messageAttributeName: string
145155
sessionId = ''
146156
dynamodbClient: DynamoDBClient
147157

148158
constructor(fields: BufferMemoryInput & BufferMemoryExtendedInput) {
149159
super(fields)
150160
this.sessionId = fields.sessionId
151161
this.dynamodbClient = fields.dynamodbClient
162+
this.tableName = fields.tableName
163+
this.partitionKey = fields.partitionKey
164+
this.dynamoKey = fields.dynamoKey
152165
}
153166

154167
overrideDynamoKey(overrideSessionId = '') {
155-
const existingDynamoKey = (this as any).dynamoKey
156-
const partitionKey = (this as any).partitionKey
168+
const existingDynamoKey = this.dynamoKey
169+
const partitionKey = this.partitionKey
157170

158171
let newDynamoKey: Record<string, AttributeValue> = {}
159172

@@ -209,9 +222,9 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
209222
async getChatMessages(overrideSessionId = '', returnBaseMessages = false): Promise<IMessage[] | BaseMessage[]> {
210223
if (!this.dynamodbClient) return []
211224

212-
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : (this as any).dynamoKey
213-
const tableName = (this as any).tableName
214-
const messageAttributeName = (this as any).messageAttributeName
225+
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : this.dynamoKey
226+
const tableName = this.tableName
227+
const messageAttributeName = this.messageAttributeName
215228

216229
const params: GetItemCommandInput = {
217230
TableName: tableName,
@@ -236,9 +249,9 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
236249
async addChatMessages(msgArray: { text: string; type: MessageType }[], overrideSessionId = ''): Promise<void> {
237250
if (!this.dynamodbClient) return
238251

239-
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : (this as any).dynamoKey
240-
const tableName = (this as any).tableName
241-
const messageAttributeName = (this as any).messageAttributeName
252+
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : this.dynamoKey
253+
const tableName = this.tableName
254+
const messageAttributeName = this.messageAttributeName
242255

243256
const input = msgArray.find((msg) => msg.type === 'userMessage')
244257
const output = msgArray.find((msg) => msg.type === 'apiMessage')
@@ -259,8 +272,8 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
259272
async clearChatMessages(overrideSessionId = ''): Promise<void> {
260273
if (!this.dynamodbClient) return
261274

262-
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : (this as any).dynamoKey
263-
const tableName = (this as any).tableName
275+
const dynamoKey = overrideSessionId ? this.overrideDynamoKey(overrideSessionId) : this.dynamoKey
276+
const tableName = this.tableName
264277

265278
const params: DeleteItemCommandInput = {
266279
TableName: tableName,

0 commit comments

Comments
 (0)