add top K to vector stores

This commit is contained in:
Henry
2023-06-10 17:21:27 +01:00
parent 617b3bea96
commit d11cb5f4b4
36 changed files with 925 additions and 238 deletions
@@ -35,6 +35,14 @@ class InMemoryVectorStore_VectorStores implements INode {
label: 'Embeddings',
name: 'embeddings',
type: 'Embeddings'
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
optional: true
}
]
this.outputs = [
@@ -55,6 +63,8 @@ class InMemoryVectorStore_VectorStores implements INode {
const docs = nodeData.inputs?.document as Document[]
const embeddings = nodeData.inputs?.embeddings as Embeddings
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const flattenDocs = docs && docs.length ? flatten(docs) : []
const finalDocs = []
@@ -65,9 +75,10 @@ class InMemoryVectorStore_VectorStores implements INode {
const vectorStore = await MemoryVectorStore.fromDocuments(finalDocs, embeddings)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore