mirror of
https://github.com/farcasclaudiu/Flowise.git
synced 2026-06-28 13:00:56 +03:00
add top K to vector stores
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user