Feature/MultiRetrievalQAChain

packages/components/nodes/chains/MultiRetrievalQAChain/MultiRetrievalQAChain.ts

@@ -0,0 +1,68 @@
+import { BaseLanguageModel } from 'langchain/base_language'
+import { INode, INodeData, INodeParams, VectorStoreRetriever } from '../../../src/Interface'
+import { getBaseClasses } from '../../../src/utils'
+import { MultiRetrievalQAChain } from 'langchain/chains'
+
+class MultiRetrievalQAChain_Chains implements INode {
+    label: string
+    name: string
+    type: string
+    icon: string
+    category: string
+    baseClasses: string[]
+    description: string
+    inputs: INodeParams[]
+
+    constructor() {
+        this.label = 'Multi Retrieval QA Chain'
+        this.name = 'multiRetrievalQAChain'
+        this.type = 'MultiRetrievalQAChain'
+        this.icon = 'chain.svg'
+        this.category = 'Chains'
+        this.description = 'QA Chain that automatically picks an appropriate vector store from multiple retrievers'
+        this.baseClasses = [this.type, ...getBaseClasses(MultiRetrievalQAChain)]
+        this.inputs = [
+            {
+                label: 'Language Model',
+                name: 'model',
+                type: 'BaseLanguageModel'
+            },
+            {
+                label: 'Vector Store Retriever',
+                name: 'vectorStoreRetriever',
+                type: 'VectorStoreRetriever',
+                list: true
+            }
+        ]
+    }
+
+    async init(nodeData: INodeData): Promise<any> {
+        const model = nodeData.inputs?.model as BaseLanguageModel
+        const vectorStoreRetriever = nodeData.inputs?.vectorStoreRetriever as VectorStoreRetriever[]
+        const retrieverNames = []
+        const retrieverDescriptions = []
+        const retrievers = []
+
+        for (const vs of vectorStoreRetriever) {
+            retrieverNames.push(vs.name)
+            retrieverDescriptions.push(vs.description)
+            retrievers.push(vs.vectorStore.asRetriever())
+        }
+
+        const chain = MultiRetrievalQAChain.fromRetrievers(model, retrieverNames, retrieverDescriptions, retrievers, undefined, {
+            verbose: process.env.DEBUG === 'true' ? true : false
+        } as any)
+
+        return chain
+    }
+
+    async run(nodeData: INodeData, input: string): Promise<string> {
+        const chain = nodeData.instance as MultiRetrievalQAChain
+
+        const res = await chain.call({ input })
+
+        return res?.text
+    }
+}
+
+module.exports = { nodeClass: MultiRetrievalQAChain_Chains }

packages/components/nodes/retrievers/VectorStoreRetriever/VectorStoreRetriever.ts

@@ -0,0 +1,61 @@
+import { VectorStore } from 'langchain/vectorstores/base'
+import { INode, INodeData, INodeParams, VectorStoreRetriever, VectorStoreRetrieverInput } from '../../../src/Interface'
+
+class VectorStoreRetriever_Retrievers implements INode {
+    label: string
+    name: string
+    description: string
+    type: string
+    icon: string
+    category: string
+    baseClasses: string[]
+    inputs: INodeParams[]
+
+    constructor() {
+        this.label = 'Vector Store Retriever'
+        this.name = 'vectorStoreRetriever'
+        this.type = 'VectorStoreRetriever'
+        this.icon = 'vectorretriever.svg'
+        this.category = 'Retrievers'
+        this.description = 'Store vector store as retriever to be later queried by MultiRetrievalQAChain'
+        this.baseClasses = [this.type]
+        this.inputs = [
+            {
+                label: 'Vector Store',
+                name: 'vectorStore',
+                type: 'VectorStore'
+            },
+            {
+                label: 'Retriever Name',
+                name: 'name',
+                type: 'string',
+                placeholder: 'netflix movies'
+            },
+            {
+                label: 'Retriever Description',
+                name: 'description',
+                type: 'string',
+                rows: 3,
+                description: 'Description of when to use the vector store retriever',
+                placeholder: 'Good for answering questions about netflix movies'
+            }
+        ]
+    }
+
+    async init(nodeData: INodeData): Promise<any> {
+        const name = nodeData.inputs?.name as string
+        const description = nodeData.inputs?.description as string
+        const vectorStore = nodeData.inputs?.vectorStore as VectorStore
+
+        const obj = {
+            name,
+            description,
+            vectorStore
+        } as VectorStoreRetrieverInput
+
+        const retriever = new VectorStoreRetriever(obj)
+        return retriever
+    }
+}
+
+module.exports = { nodeClass: VectorStoreRetriever_Retrievers }

packages/components/src/Interface.ts

@@ -95,6 +95,7 @@ export interface IMessage {
  */
 
 import { PromptTemplate as LangchainPromptTemplate, PromptTemplateInput } from 'langchain/prompts'
+import { VectorStore } from 'langchain/vectorstores/base'
 
 export class PromptTemplate extends LangchainPromptTemplate {
     promptValues: ICommonObject
@@ -124,3 +125,21 @@ export class PromptRetriever {
         this.systemMessage = `${fields.systemMessage}\n${fixedTemplate}`
     }
 }
+
+export interface VectorStoreRetrieverInput {
+    name: string
+    description: string
+    vectorStore: VectorStore
+}
+
+export class VectorStoreRetriever {
+    name: string
+    description: string
+    vectorStore: VectorStore
+
+    constructor(fields: VectorStoreRetrieverInput) {
+        this.name = fields.name
+        this.description = fields.description
+        this.vectorStore = fields.vectorStore
+    }
+}