This commit is contained in:
Ashes47
2024-01-28 21:39:22 +05:30
parent e154461f1d
commit 36ab1681ac
@@ -28,20 +28,20 @@ class MongoDB_Memory implements INode {
inputs: INodeParams[] inputs: INodeParams[]
constructor() { constructor() {
this.label = 'MongoDB Atlas Chat Memory'; this.label = 'MongoDB Atlas Chat Memory'
this.name = 'MongoDBAtlasChatMemory'; this.name = 'MongoDBAtlasChatMemory'
this.version = 1.0; this.version = 1.0
this.type = 'MongoDBAtlasChatMemory'; this.type = 'MongoDBAtlasChatMemory'
this.icon = 'mongodb.svg'; this.icon = 'mongodb.svg'
this.category = 'Memory'; this.category = 'Memory'
this.description = 'Stores the conversation in MongoDB Atlas'; this.description = 'Stores the conversation in MongoDB Atlas'
this.baseClasses = [this.type, ...getBaseClasses(BufferMemory)]; this.baseClasses = [this.type, ...getBaseClasses(BufferMemory)]
this.credential = { this.credential = {
label: 'Connect Credential', label: 'Connect Credential',
name: 'credential', name: 'credential',
type: 'credential', type: 'credential',
credentialNames: ['mongoDBUrlApi'] credentialNames: ['mongoDBUrlApi']
}; }
this.inputs = [ this.inputs = [
{ {
label: 'Database', label: 'Database',
@@ -59,7 +59,8 @@ class MongoDB_Memory implements INode {
label: 'Session Id', label: 'Session Id',
name: 'sessionId', name: 'sessionId',
type: 'string', type: 'string',
description: 'If not specified, a random id will be used. Learn <a target="_blank" href="https://docs.flowiseai.com/memory/long-term-memory#ui-and-embedded-chat">more</a>', description:
'If not specified, a random id will be used. Learn <a target="_blank" href="https://docs.flowiseai.com/memory/long-term-memory#ui-and-embedded-chat">more</a>',
default: '', default: '',
additionalParams: true, additionalParams: true,
optional: true optional: true
@@ -71,126 +72,126 @@ class MongoDB_Memory implements INode {
default: 'chat_history', default: 'chat_history',
additionalParams: true additionalParams: true
} }
]; ]
} }
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> { async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
return initializeMongoDB(nodeData, options); return initializeMongoDB(nodeData, options)
} }
} }
const initializeMongoDB = async (nodeData: INodeData, options: ICommonObject): Promise<BufferMemory> => { const initializeMongoDB = async (nodeData: INodeData, options: ICommonObject): Promise<BufferMemory> => {
const databaseName = nodeData.inputs?.databaseName as string; const databaseName = nodeData.inputs?.databaseName as string
const collectionName = nodeData.inputs?.collectionName as string; const collectionName = nodeData.inputs?.collectionName as string
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
const credentialData = await getCredentialData(nodeData.credential ?? '', options); const credentialData = await getCredentialData(nodeData.credential ?? '', options)
const mongoDBConnectUrl = getCredentialParam('mongoDBConnectUrl', credentialData, nodeData); const mongoDBConnectUrl = getCredentialParam('mongoDBConnectUrl', credentialData, nodeData)
const client = await getMongoClient(mongoDBConnectUrl); const client = await getMongoClient(mongoDBConnectUrl)
const collection = client.db(databaseName).collection(collectionName); const collection = client.db(databaseName).collection(collectionName)
const mongoDBChatMessageHistory = new MongoDBChatMessageHistory({ const mongoDBChatMessageHistory = new MongoDBChatMessageHistory({
collection, collection,
sessionId sessionId
}); })
mongoDBChatMessageHistory.getMessages = async (): Promise<BaseMessage[]> => { mongoDBChatMessageHistory.getMessages = async (): Promise<BaseMessage[]> => {
const document = await collection.findOne({ const document = await collection.findOne({
sessionId: (mongoDBChatMessageHistory as any).sessionId sessionId: (mongoDBChatMessageHistory as any).sessionId
}); })
const messages = document?.messages || []; const messages = document?.messages || []
return messages.map(mapStoredMessageToChatMessage); return messages.map(mapStoredMessageToChatMessage)
}; }
mongoDBChatMessageHistory.addMessage = async (message: BaseMessage): Promise<void> => { mongoDBChatMessageHistory.addMessage = async (message: BaseMessage): Promise<void> => {
const messages = [message].map((msg) => msg.toDict()); const messages = [message].map((msg) => msg.toDict())
await collection.updateOne( await collection.updateOne(
{ sessionId: (mongoDBChatMessageHistory as any).sessionId }, { sessionId: (mongoDBChatMessageHistory as any).sessionId },
{ {
$push: { messages: { $each: messages } } $push: { messages: { $each: messages } }
}, },
{ upsert: true } { upsert: true }
); )
}; }
mongoDBChatMessageHistory.clear = async (): Promise<void> => { mongoDBChatMessageHistory.clear = async (): Promise<void> => {
await collection.deleteOne({ sessionId: (mongoDBChatMessageHistory as any).sessionId }); await collection.deleteOne({ sessionId: (mongoDBChatMessageHistory as any).sessionId })
}; }
return new BufferMemoryExtended({ return new BufferMemoryExtended({
memoryKey: memoryKey ?? 'chat_history', memoryKey: memoryKey ?? 'chat_history',
chatHistory: mongoDBChatMessageHistory, chatHistory: mongoDBChatMessageHistory,
sessionId, sessionId,
collection collection
}); })
}; }
interface BufferMemoryExtendedInput { interface BufferMemoryExtendedInput {
collection: Collection<Document>; collection: Collection<Document>
sessionId: string; sessionId: string
} }
class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods { class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
sessionId = ''; sessionId = ''
collection: Collection<Document>; collection: Collection<Document>
constructor(fields: BufferMemoryInput & BufferMemoryExtendedInput) { constructor(fields: BufferMemoryInput & BufferMemoryExtendedInput) {
super(fields); super(fields)
this.sessionId = fields.sessionId; this.sessionId = fields.sessionId
this.collection = fields.collection; this.collection = fields.collection
} }
async getChatMessages(overrideSessionId = '', returnBaseMessages = false): Promise<IMessage[] | BaseMessage[]> { async getChatMessages(overrideSessionId = '', returnBaseMessages = false): Promise<IMessage[] | BaseMessage[]> {
if (!this.collection) return []; if (!this.collection) return []
const id = overrideSessionId ?? this.sessionId; const id = overrideSessionId ?? this.sessionId
const document = await this.collection.findOne({ sessionId: id }); const document = await this.collection.findOne({ sessionId: id })
const messages = document?.messages || []; const messages = document?.messages || []
const baseMessages = messages.map(mapStoredMessageToChatMessage); const baseMessages = messages.map(mapStoredMessageToChatMessage)
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.collection) return; if (!this.collection) return
const id = overrideSessionId ?? this.sessionId; const id = overrideSessionId ?? this.sessionId
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')
if (input) { if (input) {
const newInputMessage = new HumanMessage(input.text); const newInputMessage = new HumanMessage(input.text)
const messageToAdd = [newInputMessage].map((msg) => msg.toDict()); const messageToAdd = [newInputMessage].map((msg) => msg.toDict())
await this.collection.updateOne( await this.collection.updateOne(
{ sessionId: id }, { sessionId: id },
{ {
$push: { messages: { $each: messageToAdd } } $push: { messages: { $each: messageToAdd } }
}, },
{ upsert: true } { upsert: true }
); )
} }
if (output) { if (output) {
const newOutputMessage = new AIMessage(output.text); const newOutputMessage = new AIMessage(output.text)
const messageToAdd = [newOutputMessage].map((msg) => msg.toDict()); const messageToAdd = [newOutputMessage].map((msg) => msg.toDict())
await this.collection.updateOne( await this.collection.updateOne(
{ sessionId: id }, { sessionId: id },
{ {
$push: { messages: { $each: messageToAdd } } $push: { messages: { $each: messageToAdd } }
}, },
{ upsert: true } { upsert: true }
); )
} }
} }
async clearChatMessages(overrideSessionId = ''): Promise<void> { async clearChatMessages(overrideSessionId = ''): Promise<void> {
if (!this.collection) return; if (!this.collection) return
const id = overrideSessionId ?? this.sessionId; const id = overrideSessionId ?? this.sessionId
await this.collection.deleteOne({ sessionId: id }); await this.collection.deleteOne({ sessionId: id })
await this.clear(); await this.clear()
} }
} }
module.exports = { nodeClass: MongoDB_Memory }; module.exports = { nodeClass: MongoDB_Memory }