update all memory nodes

packages/components/nodes/memory/UpstashRedisBackedChatMemory/UpstashRedisBackedChatMemory.ts

@@ -1,8 +1,16 @@
-import { INode, INodeData, INodeParams } from '../../../src/Interface'
-import { getBaseClasses, getCredentialData, getCredentialParam, serializeChatHistory } from '../../../src/utils'
-import { ICommonObject } from '../../../src'
+import { Redis } from '@upstash/redis'
 import { BufferMemory, BufferMemoryInput } from 'langchain/memory'
 import { UpstashRedisChatMessageHistory } from 'langchain/stores/message/upstash_redis'
+import { mapStoredMessageToChatMessage, AIMessage, HumanMessage, StoredMessage, BaseMessage } from 'langchain/schema'
+import { FlowiseMemory, IMessage, INode, INodeData, INodeParams, MemoryMethods, MessageType } from '../../../src/Interface'
+import {
+    convertBaseMessagetoIMessage,
+    getBaseClasses,
+    getCredentialData,
+    getCredentialParam,
+    serializeChatHistory
+} from '../../../src/utils'
+import { ICommonObject } from '../../../src/Interface'
 
 class UpstashRedisBackedChatMemory_Memory implements INode {
     label: string
@@ -84,29 +92,39 @@ class UpstashRedisBackedChatMemory_Memory implements INode {
 
 const initalizeUpstashRedis = async (nodeData: INodeData, options: ICommonObject): Promise<BufferMemory> => {
     const baseURL = nodeData.inputs?.baseURL as string
-    const sessionId = nodeData.inputs?.sessionId as string
     const sessionTTL = nodeData.inputs?.sessionTTL as string
     const chatId = options?.chatId as string
 
     let isSessionIdUsingChatMessageId = false
-    if (!sessionId && chatId) isSessionIdUsingChatMessageId = true
+    let sessionId = ''
+
+    if (!nodeData.inputs?.sessionId && chatId) {
+        isSessionIdUsingChatMessageId = true
+        sessionId = chatId
+    } else {
+        sessionId = nodeData.inputs?.sessionId
+    }
 
     const credentialData = await getCredentialData(nodeData.credential ?? '', options)
     const upstashRestToken = getCredentialParam('upstashRestToken', credentialData, nodeData)
 
+    const client = new Redis({
+        url: baseURL,
+        token: upstashRestToken
+    })
+
     const redisChatMessageHistory = new UpstashRedisChatMessageHistory({
-        sessionId: sessionId ? sessionId : chatId,
+        sessionId,
         sessionTTL: sessionTTL ? parseInt(sessionTTL, 10) : undefined,
-        config: {
-            url: baseURL,
-            token: upstashRestToken
-        }
+        client
     })
 
     const memory = new BufferMemoryExtended({
         memoryKey: 'chat_history',
         chatHistory: redisChatMessageHistory,
-        isSessionIdUsingChatMessageId
+        isSessionIdUsingChatMessageId,
+        sessionId,
+        redisClient: client
     })
 
     return memory
@@ -114,14 +132,59 @@ const initalizeUpstashRedis = async (nodeData: INodeData, options: ICommonObject
 
 interface BufferMemoryExtendedInput {
     isSessionIdUsingChatMessageId: boolean
+    redisClient: Redis
+    sessionId: string
 }
 
-class BufferMemoryExtended extends BufferMemory {
+class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
     isSessionIdUsingChatMessageId? = false
+    sessionId = ''
+    redisClient: Redis
 
-    constructor(fields: BufferMemoryInput & Partial<BufferMemoryExtendedInput>) {
+    constructor(fields: BufferMemoryInput & BufferMemoryExtendedInput) {
         super(fields)
         this.isSessionIdUsingChatMessageId = fields.isSessionIdUsingChatMessageId
+        this.sessionId = fields.sessionId
+        this.redisClient = fields.redisClient
+    }
+
+    async getChatMessages(overrideSessionId = '', returnBaseMessages = false): Promise<IMessage[] | BaseMessage[]> {
+        if (!this.redisClient) return []
+
+        const id = overrideSessionId ?? this.sessionId
+        const rawStoredMessages: StoredMessage[] = await this.redisClient.lrange<StoredMessage>(id, 0, -1)
+        const orderedMessages = rawStoredMessages.reverse()
+        const previousMessages = orderedMessages.filter((x): x is StoredMessage => x.type !== undefined && x.data.content !== undefined)
+        const baseMessages = previousMessages.map(mapStoredMessageToChatMessage)
+        return returnBaseMessages ? baseMessages : convertBaseMessagetoIMessage(baseMessages)
+    }
+
+    async addChatMessages(msgArray: { text: string; type: MessageType }[], overrideSessionId = ''): Promise<void> {
+        if (!this.redisClient) return
+
+        const id = overrideSessionId ?? this.sessionId
+        const input = msgArray.find((msg) => msg.type === 'userMessage')
+        const output = msgArray.find((msg) => msg.type === 'apiMessage')
+
+        if (input) {
+            const newInputMessage = new HumanMessage(input.text)
+            const messageToAdd = [newInputMessage].map((msg) => msg.toDict())
+            await this.redisClient.lpush(id, JSON.stringify(messageToAdd[0]))
+        }
+
+        if (output) {
+            const newOutputMessage = new AIMessage(output.text)
+            const messageToAdd = [newOutputMessage].map((msg) => msg.toDict())
+            await this.redisClient.lpush(id, JSON.stringify(messageToAdd[0]))
+        }
+    }
+
+    async clearChatMessages(overrideSessionId = ''): Promise<void> {
+        if (!this.redisClient) return
+
+        const id = overrideSessionId ?? this.sessionId
+        await this.redisClient.del(id)
+        await this.clear()
     }
 }