update multi chains

katmakhan · Jun 14, 2023 · e50c065 · e50c065
1 parent 74939c1
commit e50c065
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 10 deletions.
diff --git a/packages/components/nodes/chains/MultiPromptChain/MultiPromptChain.ts b/packages/components/nodes/chains/MultiPromptChain/MultiPromptChain.ts
@@ -49,9 +49,12 @@ class MultiPromptChain_Chains implements INode {
             promptTemplates.push(prompt.systemMessage)
         }
 
-        const chain = MultiPromptChain.fromPrompts(model, promptNames, promptDescriptions, promptTemplates, undefined, {
-            verbose: process.env.DEBUG === 'true' ? true : false
-        } as any)
+        const chain = MultiPromptChain.fromLLMAndPrompts(model, {
+            promptNames,
+            promptDescriptions,
+            promptTemplates,
+            llmChainOpts: { verbose: process.env.DEBUG === 'true' ? true : false }
+        })
 
         return chain
     }

diff --git a/packages/components/nodes/chains/MultiRetrievalQAChain/MultiRetrievalQAChain.ts b/packages/components/nodes/chains/MultiRetrievalQAChain/MultiRetrievalQAChain.ts
@@ -32,13 +32,21 @@ class MultiRetrievalQAChain_Chains implements INode {
                 name: 'vectorStoreRetriever',
                 type: 'VectorStoreRetriever',
                 list: true
+            },
+            {
+                label: 'Return Source Documents',
+                name: 'returnSourceDocuments',
+                type: 'boolean',
+                optional: true
             }
         ]
     }
 
     async init(nodeData: INodeData): Promise<any> {
         const model = nodeData.inputs?.model as BaseLanguageModel
         const vectorStoreRetriever = nodeData.inputs?.vectorStoreRetriever as VectorStoreRetriever[]
+        const returnSourceDocuments = nodeData.inputs?.returnSourceDocuments as boolean
+
         const retrieverNames = []
         const retrieverDescriptions = []
         const retrievers = []
@@ -49,23 +57,29 @@ class MultiRetrievalQAChain_Chains implements INode {
             retrievers.push(vs.vectorStore.asRetriever((vs.vectorStore as any).k ?? 4))
         }
 
-        const chain = MultiRetrievalQAChain.fromRetrievers(model, retrieverNames, retrieverDescriptions, retrievers, undefined, {
-            verbose: process.env.DEBUG === 'true' ? true : false
-        } as any)
-
+        const chain = MultiRetrievalQAChain.fromLLMAndRetrievers(model, {
+            retrieverNames,
+            retrieverDescriptions,
+            retrievers,
+            retrievalQAChainOpts: { verbose: process.env.DEBUG === 'true' ? true : false, returnSourceDocuments }
+        })
         return chain
     }
 
-    async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> {
+    async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string | ICommonObject> {
         const chain = nodeData.instance as MultiRetrievalQAChain
+        const returnSourceDocuments = nodeData.inputs?.returnSourceDocuments as boolean
+
         const obj = { input }
 
         if (options.socketIO && options.socketIOClientId) {
-            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId, 2)
+            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId, 2, returnSourceDocuments)
             const res = await chain.call(obj, [handler])
+            if (res.text && res.sourceDocuments) return res
             return res?.text
         } else {
             const res = await chain.call(obj)
+            if (res.text && res.sourceDocuments) return res
             return res?.text
         }
     }

diff --git a/packages/server/marketplaces/Multi Retrieval QA Chain.json b/packages/server/marketplaces/Multi Retrieval QA Chain.json
@@ -84,7 +84,14 @@
                 "baseClasses": ["MultiRetrievalQAChain", "MultiRouteChain", "BaseChain", "BaseLangChain"],
                 "category": "Chains",
                 "description": "QA Chain that automatically picks an appropriate vector store from multiple retrievers",
-                "inputParams": [],
+                "inputParams": [
+                    {
+                        "label": "Return Source Documents",
+                        "name": "returnSourceDocuments",
+                        "type": "boolean",
+                        "optional": true
+                    }
+                ],
                 "inputAnchors": [
                     {
                         "label": "Language Model",