Enable Filters in MongoDB Vector Search Retriever (#3333)

* added: support to metadata filter using mongodb vector search retriever

* fix: format lint
This commit is contained in:
Anthony Bryan Gavilan Vinces
2024-10-16 10:35:37 -05:00
committed by GitHub
parent 6c3541915b
commit 235fcfeb40
@@ -86,6 +86,13 @@ class MongoDBAtlas_VectorStores implements INode {
additionalParams: true, additionalParams: true,
optional: true optional: true
}, },
{
label: 'Mongodb Metadata Filter',
name: 'mongoMetadataFilter',
type: 'json',
optional: true,
additionalParams: true
},
{ {
label: 'Top K', label: 'Top K',
name: 'topK', name: 'topK',
@@ -164,9 +171,12 @@ class MongoDBAtlas_VectorStores implements INode {
let textKey = nodeData.inputs?.textKey as string let textKey = nodeData.inputs?.textKey as string
let embeddingKey = nodeData.inputs?.embeddingKey as string let embeddingKey = nodeData.inputs?.embeddingKey as string
const embeddings = nodeData.inputs?.embeddings as Embeddings const embeddings = nodeData.inputs?.embeddings as Embeddings
const mongoMetadataFilter = nodeData.inputs?.mongoMetadataFilter as object
let mongoDBConnectUrl = getCredentialParam('mongoDBConnectUrl', credentialData, nodeData) let mongoDBConnectUrl = getCredentialParam('mongoDBConnectUrl', credentialData, nodeData)
const filter: MongoDBAtlasVectorSearch['FilterType'] = {}
const mongoClient = await getMongoClient(mongoDBConnectUrl) const mongoClient = await getMongoClient(mongoDBConnectUrl)
try { try {
const collection = mongoClient.db(databaseName).collection(collectionName) const collection = mongoClient.db(databaseName).collection(collectionName)
@@ -181,7 +191,20 @@ class MongoDBAtlas_VectorStores implements INode {
embeddingKey embeddingKey
}) as unknown as VectorStore }) as unknown as VectorStore
return resolveVectorStoreOrRetriever(nodeData, vectorStore) if (mongoMetadataFilter) {
const metadataFilter = typeof mongoMetadataFilter === 'object' ? mongoMetadataFilter : JSON.parse(mongoMetadataFilter)
for (const key in metadataFilter) {
filter.preFilter = {
...filter.preFilter,
[key]: {
$eq: metadataFilter[key]
}
}
}
}
return resolveVectorStoreOrRetriever(nodeData, vectorStore, filter)
} catch (e) { } catch (e) {
throw new Error(e) throw new Error(e)
} }