60 lines
2.1 KiB
TypeScript
60 lines
2.1 KiB
TypeScript
import { flatten } from 'lodash'
|
|
import { Collection } from 'mongodb'
|
|
import { Embeddings } from 'langchain/embeddings/base'
|
|
import { Document } from 'langchain/document'
|
|
import { VectorStore } from 'langchain/vectorstores/base'
|
|
import { MongoDBAtlasVectorSearch } from 'langchain/vectorstores/mongodb_atlas'
|
|
import { ICommonObject, INode, INodeData } from '../../../src/Interface'
|
|
import { MongoDBSearchBase } from './MongoDBSearchBase'
|
|
|
|
class MongoDBUpsert_VectorStores extends MongoDBSearchBase implements INode {
|
|
constructor() {
|
|
super()
|
|
this.label = 'MongoDB Atlas Upsert Document'
|
|
this.name = 'MongoDBUpsert'
|
|
this.version = 1.0
|
|
this.description = 'Upsert documents to MongoDB Atlas'
|
|
this.inputs.unshift({
|
|
label: 'Document',
|
|
name: 'document',
|
|
type: 'Document',
|
|
list: true
|
|
})
|
|
}
|
|
|
|
async constructVectorStore(
|
|
embeddings: Embeddings,
|
|
collection: Collection,
|
|
indexName: string,
|
|
textKey: string,
|
|
embeddingKey: string,
|
|
docs: Document<Record<string, any>>[]
|
|
): Promise<VectorStore> {
|
|
const mongoDBAtlasVectorSearch = new MongoDBAtlasVectorSearch(embeddings, {
|
|
collection: collection,
|
|
indexName: indexName,
|
|
textKey: textKey,
|
|
embeddingKey: embeddingKey
|
|
})
|
|
await mongoDBAtlasVectorSearch.addDocuments(docs)
|
|
return mongoDBAtlasVectorSearch
|
|
}
|
|
|
|
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
|
|
const docs = nodeData.inputs?.document as Document[]
|
|
|
|
const flattenDocs = docs && docs.length ? flatten(docs) : []
|
|
const finalDocs = []
|
|
for (let i = 0; i < flattenDocs.length; i += 1) {
|
|
if (flattenDocs[i] && flattenDocs[i].pageContent) {
|
|
const document = new Document(flattenDocs[i])
|
|
finalDocs.push(document)
|
|
}
|
|
}
|
|
|
|
return super.init(nodeData, _, options, finalDocs)
|
|
}
|
|
}
|
|
|
|
module.exports = { nodeClass: MongoDBUpsert_VectorStores }
|