Chore/MongoDB Connection (#3469)

getting rid of singleton design and properly close connection after every interaction
This commit is contained in:
Henry Heng
2024-11-07 19:05:00 +00:00
committed by GitHub
parent a6183aba7a
commit eeb1d17f50
3 changed files with 287 additions and 102 deletions
@@ -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
@@ -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)
@@ -175,28 +172,25 @@ 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]
}
@@ -204,31 +198,11 @@ class MongoDBAtlas_VectorStores implements INode {
}
}
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 }