76 lines
2.5 KiB
TypeScript
76 lines
2.5 KiB
TypeScript
import { INodeData } from '../../src'
|
|
|
|
export const resolveVectorStoreOrRetriever = (nodeData: INodeData, vectorStore: any) => {
|
|
const output = nodeData.outputs?.output as string
|
|
const searchType = nodeData.outputs?.searchType as string
|
|
const topK = nodeData.inputs?.topK as string
|
|
const k = topK ? parseFloat(topK) : 4
|
|
|
|
if (output === 'retriever') {
|
|
if ('mmr' === searchType) {
|
|
const fetchK = nodeData.inputs?.fetchK as string
|
|
const lambda = nodeData.inputs?.lambda as string
|
|
const f = fetchK ? parseInt(fetchK) : 20
|
|
const l = lambda ? parseFloat(lambda) : 0.5
|
|
return vectorStore.asRetriever({
|
|
searchType: 'mmr',
|
|
k: k,
|
|
searchKwargs: {
|
|
fetchK: f,
|
|
lambda: l
|
|
}
|
|
})
|
|
} else {
|
|
// "searchType" is "similarity"
|
|
return vectorStore.asRetriever(k)
|
|
}
|
|
} else if (output === 'vectorStore') {
|
|
;(vectorStore as any).k = k
|
|
return vectorStore
|
|
}
|
|
}
|
|
|
|
export const addMMRInputParams = (inputs: any[]) => {
|
|
const mmrInputParams = [
|
|
{
|
|
label: 'Search Type',
|
|
name: 'searchType',
|
|
type: 'options',
|
|
default: 'similarity',
|
|
options: [
|
|
{
|
|
label: 'Similarity',
|
|
name: 'similarity'
|
|
},
|
|
{
|
|
label: 'Max Marginal Relevance',
|
|
name: 'mmr'
|
|
}
|
|
],
|
|
additionalParams: true,
|
|
optional: true
|
|
},
|
|
{
|
|
label: 'Fetch K (for MMR Search)',
|
|
name: 'fetchK',
|
|
description: 'Number of initial documents to fetch for MMR reranking. Default to 20. Used only when the search type is MMR',
|
|
placeholder: '20',
|
|
type: 'number',
|
|
additionalParams: true,
|
|
optional: true
|
|
},
|
|
{
|
|
label: 'Lambda (for MMR Search)',
|
|
name: 'lambda',
|
|
description:
|
|
'Number between 0 and 1 that determines the degree of diversity among the results, where 0 corresponds to maximum diversity and 1 to minimum diversity. Used only when the search type is MMR',
|
|
placeholder: '0.5',
|
|
type: 'number',
|
|
additionalParams: true,
|
|
optional: true
|
|
}
|
|
]
|
|
|
|
inputs.push(...mmrInputParams)
|
|
}
|