farcasclaudiu/Flowise
1
0
Fork 0
mirror of https://github.com/farcasclaudiu/Flowise.git synced 2026-06-28 23:01:09 +03:00
Code Issues Packages Projects Releases Wiki Activity

Update MongoDBMemory.ts

Browse Source
This commit is contained in:
Anuj Verma
2024-01-28 19:45:19 +05:30
committed by GitHub
parent b382dd4c43
commit e8deeb25cf
1 changed files with 76 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/components/nodes/memory/MongoDBMemory/MongoDBMemory.ts
+76 -69
View File
@@ -1,9 +1,19 @@
import { MongoClient, Collection, Document } from 'mongodb' import { MongoClient, Collection, Document } from 'mongodb';
import { MongoDBChatMessageHistory } from 'langchain/stores/message/mongodb' import { MongoDBChatMessageHistory } from 'langchain/stores/message/mongodb';
import { BufferMemory, BufferMemoryInput } from 'langchain/memory' import { BufferMemory, BufferMemoryInput } from 'langchain/memory';
import { mapStoredMessageToChatMessage, AIMessage, HumanMessage, BaseMessage } from 'langchain/schema' import { mapStoredMessageToChatMessage, AIMessage, HumanMessage, BaseMessage } from 'langchain/schema';
import { convertBaseMessagetoIMessage, getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils' import { convertBaseMessagetoIMessage, getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils';
import { FlowiseMemory, ICommonObject, IMessage, INode, INodeData, INodeParams, MemoryMethods, MessageType } from '../../../src/Interface' import { FlowiseMemory, ICommonObject, IMessage, INode, INodeData, INodeParams, MemoryMethods, MessageType } from '../../../src/Interface';
let mongoClientSingleton = null;
const getMongoClient = async (mongoDBConnectUrl) => {
if (!mongoClientSingleton) {
mongoClientSingleton = new MongoClient(mongoDBConnectUrl, { useNewUrlParser: true, useUnifiedTopology: true });
await mongoClientSingleton.connect();
}
return mongoClientSingleton;
};
class MongoDB_Memory implements INode { class MongoDB_Memory implements INode {
label: string label: string
@@ -18,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',
@@ -49,8 +59,7 @@ class MongoDB_Memory implements INode {
label: 'Session Id', label: 'Session Id',
name: 'sessionId', name: 'sessionId',
type: 'string', type: 'string',
description: 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>',
'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
@@ -62,128 +71,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 = new MongoClient(mongoDBConnectUrl) const client = await getMongoClient(mongoDBConnectUrl);
await client.connect() 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 };