Bugfix/redis connection is closed (#3591)

* Added redis open connection if it is closed

* Removed unecessary modification

* Added check connection in all methods

* Renamed method

* added await on method call

* Refactor Redis connection handling: remove singleton pattern, ensure connections are opened and closed per operation.

---------

Co-authored-by: Maicon Matsubara <maicon@fullwise.com.br>
This commit is contained in:
Maicon Matsubara 2024-12-03 13:14:40 -03:00 committed by GitHub
parent 3b804d7777
commit d20a970a7b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 57 additions and 113 deletions

View File

@ -1,7 +1,5 @@
import { Redis, RedisOptions } from 'ioredis' import { Redis, RedisOptions } from 'ioredis'
import { isEqual } from 'lodash'
import { BufferMemory, BufferMemoryInput } from 'langchain/memory' import { BufferMemory, BufferMemoryInput } from 'langchain/memory'
import { RedisChatMessageHistory, RedisChatMessageHistoryInput } from '@langchain/community/stores/message/ioredis'
import { mapStoredMessageToChatMessage, BaseMessage, AIMessage, HumanMessage } from '@langchain/core/messages' import { mapStoredMessageToChatMessage, BaseMessage, AIMessage, HumanMessage } from '@langchain/core/messages'
import { INode, INodeData, INodeParams, ICommonObject, MessageType, IMessage, MemoryMethods, FlowiseMemory } from '../../../src/Interface' import { INode, INodeData, INodeParams, ICommonObject, MessageType, IMessage, MemoryMethods, FlowiseMemory } from '../../../src/Interface'
import { import {
@ -12,42 +10,6 @@ import {
mapChatMessageToBaseMessage mapChatMessageToBaseMessage
} from '../../../src/utils' } from '../../../src/utils'
let redisClientSingleton: Redis
let redisClientOption: RedisOptions
let redisClientUrl: string
const getRedisClientbyOption = (option: RedisOptions) => {
if (!redisClientSingleton) {
// if client doesn't exists
redisClientSingleton = new Redis(option)
redisClientOption = option
return redisClientSingleton
} else if (redisClientSingleton && !isEqual(option, redisClientOption)) {
// if client exists but option changed
redisClientSingleton.quit()
redisClientSingleton = new Redis(option)
redisClientOption = option
return redisClientSingleton
}
return redisClientSingleton
}
const getRedisClientbyUrl = (url: string) => {
if (!redisClientSingleton) {
// if client doesn't exists
redisClientSingleton = new Redis(url)
redisClientUrl = url
return redisClientSingleton
} else if (redisClientSingleton && url !== redisClientUrl) {
// if client exists but option changed
redisClientSingleton.quit()
redisClientSingleton = new Redis(url)
redisClientUrl = url
return redisClientSingleton
}
return redisClientSingleton
}
class RedisBackedChatMemory_Memory implements INode { class RedisBackedChatMemory_Memory implements INode {
label: string label: string
name: string name: string
@ -114,11 +76,11 @@ class RedisBackedChatMemory_Memory implements INode {
} }
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> { async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
return await initalizeRedis(nodeData, options) return await initializeRedis(nodeData, options)
} }
} }
const initalizeRedis = async (nodeData: INodeData, options: ICommonObject): Promise<BufferMemory> => { const initializeRedis = async (nodeData: INodeData, options: ICommonObject): Promise<BufferMemory> => {
const sessionTTL = nodeData.inputs?.sessionTTL as number const sessionTTL = nodeData.inputs?.sessionTTL as number
const memoryKey = nodeData.inputs?.memoryKey as string const memoryKey = nodeData.inputs?.memoryKey as string
const sessionId = nodeData.inputs?.sessionId as string const sessionId = nodeData.inputs?.sessionId as string
@ -127,73 +89,55 @@ const initalizeRedis = async (nodeData: INodeData, options: ICommonObject): Prom
const credentialData = await getCredentialData(nodeData.credential ?? '', options) const credentialData = await getCredentialData(nodeData.credential ?? '', options)
const redisUrl = getCredentialParam('redisUrl', credentialData, nodeData) const redisUrl = getCredentialParam('redisUrl', credentialData, nodeData)
let client: Redis const redisOptions = redisUrl
? redisUrl
if (!redisUrl || redisUrl === '') { : ({
const username = getCredentialParam('redisCacheUser', credentialData, nodeData) port: parseInt(getCredentialParam('redisCachePort', credentialData, nodeData) || '6379'),
const password = getCredentialParam('redisCachePwd', credentialData, nodeData) host: getCredentialParam('redisCacheHost', credentialData, nodeData),
const portStr = getCredentialParam('redisCachePort', credentialData, nodeData) username: getCredentialParam('redisCacheUser', credentialData, nodeData),
const host = getCredentialParam('redisCacheHost', credentialData, nodeData) password: getCredentialParam('redisCachePwd', credentialData, nodeData),
const sslEnabled = getCredentialParam('redisCacheSslEnabled', credentialData, nodeData) tls: getCredentialParam('redisCacheSslEnabled', credentialData, nodeData) ? { rejectUnauthorized: false } : undefined
} as RedisOptions)
const tlsOptions = sslEnabled === true ? { tls: { rejectUnauthorized: false } } : {}
client = getRedisClientbyOption({
port: portStr ? parseInt(portStr) : 6379,
host,
username,
password,
...tlsOptions
})
} else {
client = getRedisClientbyUrl(redisUrl)
}
let obj: RedisChatMessageHistoryInput = {
sessionId,
client
}
if (sessionTTL) {
obj = {
...obj,
sessionTTL
}
}
const redisChatMessageHistory = new RedisChatMessageHistory(obj)
const memory = new BufferMemoryExtended({ const memory = new BufferMemoryExtended({
memoryKey: memoryKey ?? 'chat_history', memoryKey: memoryKey ?? 'chat_history',
chatHistory: redisChatMessageHistory,
sessionId, sessionId,
windowSize, windowSize,
redisClient: client, sessionTTL,
sessionTTL redisOptions
}) })
return memory return memory
} }
interface BufferMemoryExtendedInput { interface BufferMemoryExtendedInput {
redisClient: Redis
sessionId: string sessionId: string
windowSize?: number windowSize?: number
sessionTTL?: number sessionTTL?: number
redisOptions: RedisOptions | string
} }
class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods { class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
sessionId = '' sessionId = ''
redisClient: Redis
windowSize?: number windowSize?: number
sessionTTL?: number sessionTTL?: number
redisOptions: RedisOptions | string
constructor(fields: BufferMemoryInput & BufferMemoryExtendedInput) { constructor(fields: BufferMemoryInput & BufferMemoryExtendedInput) {
super(fields) super(fields)
this.sessionId = fields.sessionId this.sessionId = fields.sessionId
this.redisClient = fields.redisClient
this.windowSize = fields.windowSize this.windowSize = fields.windowSize
this.sessionTTL = fields.sessionTTL this.sessionTTL = fields.sessionTTL
this.redisOptions = fields.redisOptions
}
private async withRedisClient<T>(fn: (client: Redis) => Promise<T>): Promise<T> {
const client = typeof this.redisOptions === 'string' ? new Redis(this.redisOptions) : new Redis(this.redisOptions)
try {
return await fn(client)
} finally {
await client.quit()
}
} }
async getChatMessages( async getChatMessages(
@ -201,46 +145,46 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
returnBaseMessages = false, returnBaseMessages = false,
prependMessages?: IMessage[] prependMessages?: IMessage[]
): Promise<IMessage[] | BaseMessage[]> { ): Promise<IMessage[] | BaseMessage[]> {
if (!this.redisClient) return [] return this.withRedisClient(async (client) => {
const id = overrideSessionId ? overrideSessionId : this.sessionId
const id = overrideSessionId ? overrideSessionId : this.sessionId const rawStoredMessages = await client.lrange(id, this.windowSize ? this.windowSize * -1 : 0, -1)
const rawStoredMessages = await this.redisClient.lrange(id, this.windowSize ? this.windowSize * -1 : 0, -1) const orderedMessages = rawStoredMessages.reverse().map((message) => JSON.parse(message))
const orderedMessages = rawStoredMessages.reverse().map((message) => JSON.parse(message)) const baseMessages = orderedMessages.map(mapStoredMessageToChatMessage)
const baseMessages = orderedMessages.map(mapStoredMessageToChatMessage) if (prependMessages?.length) {
if (prependMessages?.length) { baseMessages.unshift(...(await mapChatMessageToBaseMessage(prependMessages)))
baseMessages.unshift(...(await mapChatMessageToBaseMessage(prependMessages))) }
} return returnBaseMessages ? baseMessages : convertBaseMessagetoIMessage(baseMessages)
return returnBaseMessages ? baseMessages : convertBaseMessagetoIMessage(baseMessages) })
} }
async addChatMessages(msgArray: { text: string; type: MessageType }[], overrideSessionId = ''): Promise<void> { async addChatMessages(msgArray: { text: string; type: MessageType }[], overrideSessionId = ''): Promise<void> {
if (!this.redisClient) return await this.withRedisClient(async (client) => {
const id = overrideSessionId ? overrideSessionId : this.sessionId
const input = msgArray.find((msg) => msg.type === 'userMessage')
const output = msgArray.find((msg) => msg.type === 'apiMessage')
const id = overrideSessionId ? overrideSessionId : this.sessionId if (input) {
const input = msgArray.find((msg) => msg.type === 'userMessage') const newInputMessage = new HumanMessage(input.text)
const output = msgArray.find((msg) => msg.type === 'apiMessage') const messageToAdd = [newInputMessage].map((msg) => msg.toDict())
await client.lpush(id, JSON.stringify(messageToAdd[0]))
if (this.sessionTTL) await client.expire(id, this.sessionTTL)
}
if (input) { if (output) {
const newInputMessage = new HumanMessage(input.text) const newOutputMessage = new AIMessage(output.text)
const messageToAdd = [newInputMessage].map((msg) => msg.toDict()) const messageToAdd = [newOutputMessage].map((msg) => msg.toDict())
await this.redisClient.lpush(id, JSON.stringify(messageToAdd[0])) await client.lpush(id, JSON.stringify(messageToAdd[0]))
if (this.sessionTTL) await this.redisClient.expire(id, this.sessionTTL) if (this.sessionTTL) await client.expire(id, this.sessionTTL)
} }
})
if (output) {
const newOutputMessage = new AIMessage(output.text)
const messageToAdd = [newOutputMessage].map((msg) => msg.toDict())
await this.redisClient.lpush(id, JSON.stringify(messageToAdd[0]))
if (this.sessionTTL) await this.redisClient.expire(id, this.sessionTTL)
}
} }
async clearChatMessages(overrideSessionId = ''): Promise<void> { async clearChatMessages(overrideSessionId = ''): Promise<void> {
if (!this.redisClient) return await this.withRedisClient(async (client) => {
const id = overrideSessionId ? overrideSessionId : this.sessionId
const id = overrideSessionId ? overrideSessionId : this.sessionId await client.del(id)
await this.redisClient.del(id) await this.clear()
await this.clear() })
} }
} }