Update UpstashRedisBackedChatMemory.ts

code cleanup
This commit is contained in:
Henry Heng 2023-10-08 01:52:19 +01:00 committed by GitHub
parent 9bae1c33f7
commit 424e713128
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 29 deletions

View File

@ -2,8 +2,7 @@ import { INode, INodeData, INodeParams } from '../../../src/Interface'
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils' import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
import { ICommonObject } from '../../../src' import { ICommonObject } from '../../../src'
import { BufferMemory, BufferMemoryInput } from 'langchain/memory' import { BufferMemory, BufferMemoryInput } from 'langchain/memory'
import { UpstashRedisChatMessageHistory, UpstashRedisChatMessageHistoryInput } from 'langchain/stores/message/upstash_redis' import { UpstashRedisChatMessageHistory } from 'langchain/stores/message/upstash_redis'
import { Redis, RedisConfigNodejs } from '@upstash/redis'
class UpstashRedisBackedChatMemory_Memory implements INode { class UpstashRedisBackedChatMemory_Memory implements INode {
label: string label: string
@ -35,10 +34,10 @@ class UpstashRedisBackedChatMemory_Memory implements INode {
} }
this.inputs = [ this.inputs = [
{ {
label: 'Base URL', label: 'Upstash Redis REST URL',
name: 'baseURL', name: 'baseURL',
type: 'string', type: 'string',
default: 'https://<your-url>.upstash.io' placeholder: 'https://<your-url>.upstash.io'
}, },
{ {
label: 'Session Id', label: 'Session Id',
@ -56,13 +55,6 @@ class UpstashRedisBackedChatMemory_Memory implements INode {
description: 'Omit this parameter to make sessions never expire', description: 'Omit this parameter to make sessions never expire',
additionalParams: true, additionalParams: true,
optional: true optional: true
},
{
label: 'Memory Key',
name: 'memoryKey',
type: 'string',
default: 'chat_history',
additionalParams: true
} }
] ]
} }
@ -84,38 +76,29 @@ class UpstashRedisBackedChatMemory_Memory implements INode {
const initalizeUpstashRedis = async (nodeData: INodeData, options: ICommonObject): Promise<BufferMemory> => { const initalizeUpstashRedis = async (nodeData: INodeData, options: ICommonObject): Promise<BufferMemory> => {
const baseURL = nodeData.inputs?.baseURL as string const baseURL = nodeData.inputs?.baseURL as string
const sessionId = nodeData.inputs?.sessionId as string const sessionId = nodeData.inputs?.sessionId as string
const sessionTTL = nodeData.inputs?.sessionTTL as number const sessionTTL = nodeData.inputs?.sessionTTL as string
const memoryKey = nodeData.inputs?.memoryKey as string
const chatId = options?.chatId as string const chatId = options?.chatId as string
let isSessionIdUsingChatMessageId = false let isSessionIdUsingChatMessageId = false
if (!sessionId && chatId) isSessionIdUsingChatMessageId = true if (!sessionId && chatId) isSessionIdUsingChatMessageId = true
const credentialData = await getCredentialData(nodeData.credential ?? '', options) const credentialData = await getCredentialData(nodeData.credential ?? '', options)
const upstashRedisPassword = getCredentialParam('upstashRedisPassword', credentialData, nodeData) const upstashRestToken = getCredentialParam('upstashRestToken', credentialData, nodeData)
const upstashRedisConfig = { url: baseURL, token: upstashRedisPassword } as RedisConfigNodejs const redisChatMessageHistory = new UpstashRedisChatMessageHistory({
const redisClient = new Redis(upstashRedisConfig)
let obj: UpstashRedisChatMessageHistoryInput = {
sessionId: sessionId ? sessionId : chatId, sessionId: sessionId ? sessionId : chatId,
client: redisClient sessionTTL: sessionTTL ? parseInt(sessionTTL, 10) : undefined,
} config: {
url: baseURL,
if (sessionTTL) { token: upstashRestToken
obj = {
...obj,
sessionTTL
} }
} })
const redisChatMessageHistory = new UpstashRedisChatMessageHistory(obj)
const memory = new BufferMemoryExtended({ const memory = new BufferMemoryExtended({
memoryKey,
chatHistory: redisChatMessageHistory, chatHistory: redisChatMessageHistory,
returnMessages: true,
isSessionIdUsingChatMessageId isSessionIdUsingChatMessageId
}) })
return memory return memory
} }