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