Skip to content

Commit

Permalink
Chore/MongoDB Connection (FlowiseAI#3469)
Browse files Browse the repository at this point in the history 
getting rid of singleton design and properly close connection after every interaction
  • Loading branch information
@HenryHengZJ
HenryHengZJ authored Nov 7, 2024
1 parent a6183ab commit eeb1d17
Show file tree
Hide file tree
Showing 3 changed files with 288 additions and 103 deletions.
106 changes: 40 additions & 66 deletions packages/components/nodes/memory/MongoDBMemory/MongoDBMemory.ts
View file
Original file line number Diff line number Diff line change
@@ -1,34 +1,18 @@
import { MongoClient, Collection, Document } from 'mongodb'
import { MongoDBChatMessageHistory } from '@langchain/mongodb'
import { MongoClient } from 'mongodb'
import { BufferMemory, BufferMemoryInput } from 'langchain/memory'
import { mapStoredMessageToChatMessage, AIMessage, HumanMessage, BaseMessage } from '@langchain/core/messages'
import {
convertBaseMessagetoIMessage,
getBaseClasses,
getCredentialData,
getCredentialParam,
getVersion,
mapChatMessageToBaseMessage
} from '../../../src/utils'
import { FlowiseMemory, ICommonObject, IMessage, INode, INodeData, INodeParams, MemoryMethods, MessageType } from '../../../src/Interface'

let mongoClientSingleton: MongoClient
let mongoUrl: string

const getMongoClient = async (newMongoUrl: string) => {
if (!mongoClientSingleton) {
// if client does not exist
mongoClientSingleton = new MongoClient(newMongoUrl)
mongoUrl = newMongoUrl
return mongoClientSingleton
} else if (mongoClientSingleton && newMongoUrl !== mongoUrl) {
// if client exists but url changed
mongoClientSingleton.close()
mongoClientSingleton = new MongoClient(newMongoUrl)
mongoUrl = newMongoUrl
return mongoClientSingleton
}
return mongoClientSingleton
}
// TODO: Add ability to specify env variable and use singleton pattern (i.e initialize MongoDB on server and pass to component)

class MongoDB_Memory implements INode {
label: string
name: string
Expand Down Expand Up @@ -102,83 +86,68 @@ const initializeMongoDB = async (nodeData: INodeData, options: ICommonObject): P

const credentialData = await getCredentialData(nodeData.credential ?? '', options)
const mongoDBConnectUrl = getCredentialParam('mongoDBConnectUrl', credentialData, nodeData)

const client = await getMongoClient(mongoDBConnectUrl)
const collection = client.db(databaseName).collection(collectionName)

const mongoDBChatMessageHistory = new MongoDBChatMessageHistory({
collection,
sessionId
})

// @ts-ignore
mongoDBChatMessageHistory.getMessages = async (): Promise<BaseMessage[]> => {
const document = await collection.findOne({
sessionId: (mongoDBChatMessageHistory as any).sessionId
})
const messages = document?.messages || []
return messages.map(mapStoredMessageToChatMessage)
}

// @ts-ignore
mongoDBChatMessageHistory.addMessage = async (message: BaseMessage): Promise<void> => {
const messages = [message].map((msg) => msg.toDict())
await collection.updateOne(
{ sessionId: (mongoDBChatMessageHistory as any).sessionId },
{
$push: { messages: { $each: messages } }
},
{ upsert: true }
)
}

mongoDBChatMessageHistory.clear = async (): Promise<void> => {
await collection.deleteOne({ sessionId: (mongoDBChatMessageHistory as any).sessionId })
}
const driverInfo = { name: 'Flowise', version: (await getVersion()).version }

return new BufferMemoryExtended({
memoryKey: memoryKey ?? 'chat_history',
// @ts-ignore
chatHistory: mongoDBChatMessageHistory,
sessionId,
collection
mongoConnection: {
databaseName,
collectionName,
mongoDBConnectUrl,
driverInfo
}
})
}

interface BufferMemoryExtendedInput {
collection: Collection<Document>
sessionId: string
mongoConnection: {
databaseName: string
collectionName: string
mongoDBConnectUrl: string
driverInfo: { name: string; version: string }
}
}

class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
sessionId = ''
collection: Collection<Document>
mongoConnection: {
databaseName: string
collectionName: string
mongoDBConnectUrl: string
driverInfo: { name: string; version: string }
}

constructor(fields: BufferMemoryInput & BufferMemoryExtendedInput) {
super(fields)
this.sessionId = fields.sessionId
this.collection = fields.collection
this.mongoConnection = fields.mongoConnection
}

async getChatMessages(
overrideSessionId = '',
returnBaseMessages = false,
prependMessages?: IMessage[]
): Promise<IMessage[] | BaseMessage[]> {
if (!this.collection) return []
const client = new MongoClient(this.mongoConnection.mongoDBConnectUrl, { driverInfo: this.mongoConnection.driverInfo })
const collection = client.db(this.mongoConnection.databaseName).collection(this.mongoConnection.collectionName)

const id = overrideSessionId ? overrideSessionId : this.sessionId
const document = await this.collection.findOne({ sessionId: id })
const document = await collection.findOne({ sessionId: id })
const messages = document?.messages || []
const baseMessages = messages.map(mapStoredMessageToChatMessage)
if (prependMessages?.length) {
baseMessages.unshift(...(await mapChatMessageToBaseMessage(prependMessages)))
}

await client.close()
return returnBaseMessages ? baseMessages : convertBaseMessagetoIMessage(baseMessages)
}

async addChatMessages(msgArray: { text: string; type: MessageType }[], overrideSessionId = ''): Promise<void> {
if (!this.collection) return
const client = new MongoClient(this.mongoConnection.mongoDBConnectUrl, { driverInfo: this.mongoConnection.driverInfo })
const collection = client.db(this.mongoConnection.databaseName).collection(this.mongoConnection.collectionName)

const id = overrideSessionId ? overrideSessionId : this.sessionId
const input = msgArray.find((msg) => msg.type === 'userMessage')
Expand All @@ -187,7 +156,7 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
if (input) {
const newInputMessage = new HumanMessage(input.text)
const messageToAdd = [newInputMessage].map((msg) => msg.toDict())
await this.collection.updateOne(
await collection.updateOne(
{ sessionId: id },
{
$push: { messages: { $each: messageToAdd } }
Expand All @@ -199,22 +168,27 @@ class BufferMemoryExtended extends FlowiseMemory implements MemoryMethods {
if (output) {
const newOutputMessage = new AIMessage(output.text)
const messageToAdd = [newOutputMessage].map((msg) => msg.toDict())
await this.collection.updateOne(
await collection.updateOne(
{ sessionId: id },
{
$push: { messages: { $each: messageToAdd } }
},
{ upsert: true }
)
}

await client.close()
}

async clearChatMessages(overrideSessionId = ''): Promise<void> {
if (!this.collection) return
const client = new MongoClient(this.mongoConnection.mongoDBConnectUrl, { driverInfo: this.mongoConnection.driverInfo })
const collection = client.db(this.mongoConnection.databaseName).collection(this.mongoConnection.collectionName)

const id = overrideSessionId ? overrideSessionId : this.sessionId
await this.collection.deleteOne({ sessionId: id })
await collection.deleteOne({ sessionId: id })
await this.clear()

await client.close()
}
}

Expand Down
48 changes: 11 additions & 37 deletions packages/components/nodes/vectorstores/MongoDBAtlas/MongoDBAtlas.ts
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,12 @@
import { flatten } from 'lodash'
import { MongoClient } from 'mongodb'
import { MongoDBAtlasVectorSearch } from '@langchain/mongodb'
import { Embeddings } from '@langchain/core/embeddings'
import { Document } from '@langchain/core/documents'
import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams, IndexingResult } from '../../../src/Interface'
import { getBaseClasses, getCredentialData, getCredentialParam, getVersion } from '../../../src/utils'
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
import { addMMRInputParams, resolveVectorStoreOrRetriever } from '../VectorStoreUtils'
import { VectorStore } from '@langchain/core/vectorstores'
import { MongoDBAtlasVectorSearch } from './core'

// TODO: Add ability to specify env variable and use singleton pattern (i.e initialize MongoDB on server and pass to component)
class MongoDBAtlas_VectorStores implements INode {
label: string
name: string
Expand Down Expand Up @@ -142,20 +141,18 @@ class MongoDBAtlas_VectorStores implements INode {
}
}

const mongoClient = await getMongoClient(mongoDBConnectUrl)
try {
const collection = mongoClient.db(databaseName).collection(collectionName)

if (!textKey || textKey === '') textKey = 'text'
if (!embeddingKey || embeddingKey === '') embeddingKey = 'embedding'

const mongoDBAtlasVectorSearch = new MongoDBAtlasVectorSearch(embeddings, {
collection,
connectionDetails: { mongoDBConnectUrl, databaseName, collectionName },
indexName,
textKey,
embeddingKey
})
await mongoDBAtlasVectorSearch.addDocuments(finalDocs)

return { numAdded: finalDocs.length, addedDocs: finalDocs }
} catch (e) {
throw new Error(e)
Expand All @@ -175,60 +172,37 @@ class MongoDBAtlas_VectorStores implements INode {

let mongoDBConnectUrl = getCredentialParam('mongoDBConnectUrl', credentialData, nodeData)

const filter: MongoDBAtlasVectorSearch['FilterType'] = {}
const mongoDbFilter: MongoDBAtlasVectorSearch['FilterType'] = {}

const mongoClient = await getMongoClient(mongoDBConnectUrl)
try {
const collection = mongoClient.db(databaseName).collection(collectionName)

if (!textKey || textKey === '') textKey = 'text'
if (!embeddingKey || embeddingKey === '') embeddingKey = 'embedding'

const vectorStore = new MongoDBAtlasVectorSearch(embeddings, {
collection,
connectionDetails: { mongoDBConnectUrl, databaseName, collectionName },
indexName,
textKey,
embeddingKey
}) as unknown as VectorStore
})

if (mongoMetadataFilter) {
const metadataFilter = typeof mongoMetadataFilter === 'object' ? mongoMetadataFilter : JSON.parse(mongoMetadataFilter)

for (const key in metadataFilter) {
filter.preFilter = {
...filter.preFilter,
mongoDbFilter.preFilter = {
...mongoDbFilter.preFilter,
[key]: {
$eq: metadataFilter[key]
}
}
}
}

return resolveVectorStoreOrRetriever(nodeData, vectorStore, filter)
return resolveVectorStoreOrRetriever(nodeData, vectorStore, mongoDbFilter)
} catch (e) {
throw new Error(e)
}
}
}

let mongoClientSingleton: MongoClient
let mongoUrl: string

const getMongoClient = async (newMongoUrl: string) => {
const driverInfo = { name: 'Flowise', version: (await getVersion()).version }

if (!mongoClientSingleton) {
// if client does not exist
mongoClientSingleton = new MongoClient(newMongoUrl, { driverInfo })
mongoUrl = newMongoUrl
return mongoClientSingleton
} else if (mongoClientSingleton && newMongoUrl !== mongoUrl) {
// if client exists but url changed
mongoClientSingleton.close()
mongoClientSingleton = new MongoClient(newMongoUrl, { driverInfo })
mongoUrl = newMongoUrl
return mongoClientSingleton
}
return mongoClientSingleton
}
module.exports = { nodeClass: MongoDBAtlas_VectorStores }
Loading

0 comments on commit eeb1d17

Please sign in to comment.