Merge pull request FlowiseAI#1726 from FlowiseAI/feature/LlamaIndex

Feature/Update Query Engine Tool
georggr · Feb 14, 2024 · ef6edd3 · ef6edd3
2 parents e2df5e9 + 778e024
commit ef6edd3
Show file tree

Hide file tree

Showing 5 changed files with 561 additions and 448 deletions.
diff --git a/packages/components/nodes/engine/QueryEngine/QueryEngine.ts b/packages/components/nodes/engine/QueryEngine/QueryEngine.ts
@@ -28,12 +28,12 @@ class QueryEngine_LlamaIndex implements INode {
     constructor(fields?: { sessionId?: string }) {
         this.label = 'Query Engine'
         this.name = 'queryEngine'
-        this.version = 1.0
+        this.version = 2.0
         this.type = 'QueryEngine'
         this.icon = 'query-engine.png'
         this.category = 'Engine'
         this.description = 'Simple query engine built to answer question over your data, without memory'
-        this.baseClasses = [this.type]
+        this.baseClasses = [this.type, 'BaseQueryEngine']
         this.tags = ['LlamaIndex']
         this.inputs = [
             {
@@ -59,52 +59,13 @@ class QueryEngine_LlamaIndex implements INode {
         this.sessionId = fields?.sessionId
     }
 
-    async init(): Promise<any> {
-        return null
+    async init(nodeData: INodeData): Promise<any> {
+        return prepareEngine(nodeData)
     }
 
     async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string | object> {
         const returnSourceDocuments = nodeData.inputs?.returnSourceDocuments as boolean
-        const vectorStoreRetriever = nodeData.inputs?.vectorStoreRetriever
-        const responseSynthesizerObj = nodeData.inputs?.responseSynthesizer
-
-        let queryEngine = new RetrieverQueryEngine(vectorStoreRetriever)
-
-        if (responseSynthesizerObj) {
-            if (responseSynthesizerObj.type === 'TreeSummarize') {
-                const responseSynthesizer = new ResponseSynthesizer({
-                    responseBuilder: new TreeSummarize(vectorStoreRetriever.serviceContext, responseSynthesizerObj.textQAPromptTemplate),
-                    serviceContext: vectorStoreRetriever.serviceContext
-                })
-                queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
-            } else if (responseSynthesizerObj.type === 'CompactAndRefine') {
-                const responseSynthesizer = new ResponseSynthesizer({
-                    responseBuilder: new CompactAndRefine(
-                        vectorStoreRetriever.serviceContext,
-                        responseSynthesizerObj.textQAPromptTemplate,
-                        responseSynthesizerObj.refinePromptTemplate
-                    ),
-                    serviceContext: vectorStoreRetriever.serviceContext
-                })
-                queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
-            } else if (responseSynthesizerObj.type === 'Refine') {
-                const responseSynthesizer = new ResponseSynthesizer({
-                    responseBuilder: new Refine(
-                        vectorStoreRetriever.serviceContext,
-                        responseSynthesizerObj.textQAPromptTemplate,
-                        responseSynthesizerObj.refinePromptTemplate
-                    ),
-                    serviceContext: vectorStoreRetriever.serviceContext
-                })
-                queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
-            } else if (responseSynthesizerObj.type === 'SimpleResponseBuilder') {
-                const responseSynthesizer = new ResponseSynthesizer({
-                    responseBuilder: new SimpleResponseBuilder(vectorStoreRetriever.serviceContext),
-                    serviceContext: vectorStoreRetriever.serviceContext
-                })
-                queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
-            }
-        }
+        const queryEngine = prepareEngine(nodeData)
 
         let text = ''
         let sourceDocuments: ICommonObject[] = []
@@ -140,4 +101,49 @@ class QueryEngine_LlamaIndex implements INode {
     }
 }
 
+const prepareEngine = (nodeData: INodeData) => {
+    const vectorStoreRetriever = nodeData.inputs?.vectorStoreRetriever
+    const responseSynthesizerObj = nodeData.inputs?.responseSynthesizer
+
+    let queryEngine = new RetrieverQueryEngine(vectorStoreRetriever)
+
+    if (responseSynthesizerObj) {
+        if (responseSynthesizerObj.type === 'TreeSummarize') {
+            const responseSynthesizer = new ResponseSynthesizer({
+                responseBuilder: new TreeSummarize(vectorStoreRetriever.serviceContext, responseSynthesizerObj.textQAPromptTemplate),
+                serviceContext: vectorStoreRetriever.serviceContext
+            })
+            queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
+        } else if (responseSynthesizerObj.type === 'CompactAndRefine') {
+            const responseSynthesizer = new ResponseSynthesizer({
+                responseBuilder: new CompactAndRefine(
+                    vectorStoreRetriever.serviceContext,
+                    responseSynthesizerObj.textQAPromptTemplate,
+                    responseSynthesizerObj.refinePromptTemplate
+                ),
+                serviceContext: vectorStoreRetriever.serviceContext
+            })
+            queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
+        } else if (responseSynthesizerObj.type === 'Refine') {
+            const responseSynthesizer = new ResponseSynthesizer({
+                responseBuilder: new Refine(
+                    vectorStoreRetriever.serviceContext,
+                    responseSynthesizerObj.textQAPromptTemplate,
+                    responseSynthesizerObj.refinePromptTemplate
+                ),
+                serviceContext: vectorStoreRetriever.serviceContext
+            })
+            queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
+        } else if (responseSynthesizerObj.type === 'SimpleResponseBuilder') {
+            const responseSynthesizer = new ResponseSynthesizer({
+                responseBuilder: new SimpleResponseBuilder(vectorStoreRetriever.serviceContext),
+                serviceContext: vectorStoreRetriever.serviceContext
+            })
+            queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
+        }
+    }
+
+    return queryEngine
+}
+
 module.exports = { nodeClass: QueryEngine_LlamaIndex }
diff --git a/packages/components/nodes/engine/SubQuestionQueryEngine/SubQuestionQueryEngine.ts b/packages/components/nodes/engine/SubQuestionQueryEngine/SubQuestionQueryEngine.ts
@@ -33,13 +33,13 @@ class SubQuestionQueryEngine_LlamaIndex implements INode {
     constructor(fields?: { sessionId?: string }) {
         this.label = 'Sub Question Query Engine'
         this.name = 'subQuestionQueryEngine'
-        this.version = 1.0
+        this.version = 2.0
         this.type = 'SubQuestionQueryEngine'
         this.icon = 'subQueryEngine.svg'
         this.category = 'Engine'
         this.description =
             'Breaks complex query into sub questions for each relevant data source, then gather all the intermediate reponses and synthesizes a final response'
-        this.baseClasses = [this.type]
+        this.baseClasses = [this.type, 'BaseQueryEngine']
         this.tags = ['LlamaIndex']
         this.inputs = [
             {
@@ -76,85 +76,13 @@ class SubQuestionQueryEngine_LlamaIndex implements INode {
         this.sessionId = fields?.sessionId
     }
 
-    async init(): Promise<any> {
-        return null
+    async init(nodeData: INodeData): Promise<any> {
+        return prepareEngine(nodeData)
     }
 
     async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string | object> {
         const returnSourceDocuments = nodeData.inputs?.returnSourceDocuments as boolean
-        const embeddings = nodeData.inputs?.embeddings as BaseEmbedding
-        const model = nodeData.inputs?.model
-
-        const serviceContext = serviceContextFromDefaults({
-            llm: model,
-            embedModel: embeddings
-        })
-
-        let queryEngineTools = nodeData.inputs?.queryEngineTools as QueryEngineTool[]
-        queryEngineTools = flatten(queryEngineTools)
-
-        let queryEngine = SubQuestionQueryEngine.fromDefaults({
-            serviceContext,
-            queryEngineTools,
-            questionGen: new LLMQuestionGenerator({ llm: model })
-        })
-
-        const responseSynthesizerObj = nodeData.inputs?.responseSynthesizer
-        if (responseSynthesizerObj) {
-            if (responseSynthesizerObj.type === 'TreeSummarize') {
-                const responseSynthesizer = new ResponseSynthesizer({
-                    responseBuilder: new TreeSummarize(serviceContext, responseSynthesizerObj.textQAPromptTemplate),
-                    serviceContext
-                })
-                queryEngine = SubQuestionQueryEngine.fromDefaults({
-                    responseSynthesizer,
-                    serviceContext,
-                    queryEngineTools,
-                    questionGen: new LLMQuestionGenerator({ llm: model })
-                })
-            } else if (responseSynthesizerObj.type === 'CompactAndRefine') {
-                const responseSynthesizer = new ResponseSynthesizer({
-                    responseBuilder: new CompactAndRefine(
-                        serviceContext,
-                        responseSynthesizerObj.textQAPromptTemplate,
-                        responseSynthesizerObj.refinePromptTemplate
-                    ),
-                    serviceContext
-                })
-                queryEngine = SubQuestionQueryEngine.fromDefaults({
-                    responseSynthesizer,
-                    serviceContext,
-                    queryEngineTools,
-                    questionGen: new LLMQuestionGenerator({ llm: model })
-                })
-            } else if (responseSynthesizerObj.type === 'Refine') {
-                const responseSynthesizer = new ResponseSynthesizer({
-                    responseBuilder: new Refine(
-                        serviceContext,
-                        responseSynthesizerObj.textQAPromptTemplate,
-                        responseSynthesizerObj.refinePromptTemplate
-                    ),
-                    serviceContext
-                })
-                queryEngine = SubQuestionQueryEngine.fromDefaults({
-                    responseSynthesizer,
-                    serviceContext,
-                    queryEngineTools,
-                    questionGen: new LLMQuestionGenerator({ llm: model })
-                })
-            } else if (responseSynthesizerObj.type === 'SimpleResponseBuilder') {
-                const responseSynthesizer = new ResponseSynthesizer({
-                    responseBuilder: new SimpleResponseBuilder(serviceContext),
-                    serviceContext
-                })
-                queryEngine = SubQuestionQueryEngine.fromDefaults({
-                    responseSynthesizer,
-                    serviceContext,
-                    queryEngineTools,
-                    questionGen: new LLMQuestionGenerator({ llm: model })
-                })
-            }
-        }
+        const queryEngine = prepareEngine(nodeData)
 
         let text = ''
         let sourceDocuments: ICommonObject[] = []
@@ -190,4 +118,82 @@ class SubQuestionQueryEngine_LlamaIndex implements INode {
     }
 }
 
+const prepareEngine = (nodeData: INodeData) => {
+    const embeddings = nodeData.inputs?.embeddings as BaseEmbedding
+    const model = nodeData.inputs?.model
+
+    const serviceContext = serviceContextFromDefaults({
+        llm: model,
+        embedModel: embeddings
+    })
+
+    let queryEngineTools = nodeData.inputs?.queryEngineTools as QueryEngineTool[]
+    queryEngineTools = flatten(queryEngineTools)
+
+    let queryEngine = SubQuestionQueryEngine.fromDefaults({
+        serviceContext,
+        queryEngineTools,
+        questionGen: new LLMQuestionGenerator({ llm: model })
+    })
+
+    const responseSynthesizerObj = nodeData.inputs?.responseSynthesizer
+    if (responseSynthesizerObj) {
+        if (responseSynthesizerObj.type === 'TreeSummarize') {
+            const responseSynthesizer = new ResponseSynthesizer({
+                responseBuilder: new TreeSummarize(serviceContext, responseSynthesizerObj.textQAPromptTemplate),
+                serviceContext
+            })
+            queryEngine = SubQuestionQueryEngine.fromDefaults({
+                responseSynthesizer,
+                serviceContext,
+                queryEngineTools,
+                questionGen: new LLMQuestionGenerator({ llm: model })
+            })
+        } else if (responseSynthesizerObj.type === 'CompactAndRefine') {
+            const responseSynthesizer = new ResponseSynthesizer({
+                responseBuilder: new CompactAndRefine(
+                    serviceContext,
+                    responseSynthesizerObj.textQAPromptTemplate,
+                    responseSynthesizerObj.refinePromptTemplate
+                ),
+                serviceContext
+            })
+            queryEngine = SubQuestionQueryEngine.fromDefaults({
+                responseSynthesizer,
+                serviceContext,
+                queryEngineTools,
+                questionGen: new LLMQuestionGenerator({ llm: model })
+            })
+        } else if (responseSynthesizerObj.type === 'Refine') {
+            const responseSynthesizer = new ResponseSynthesizer({
+                responseBuilder: new Refine(
+                    serviceContext,
+                    responseSynthesizerObj.textQAPromptTemplate,
+                    responseSynthesizerObj.refinePromptTemplate
+                ),
+                serviceContext
+            })
+            queryEngine = SubQuestionQueryEngine.fromDefaults({
+                responseSynthesizer,
+                serviceContext,
+                queryEngineTools,
+                questionGen: new LLMQuestionGenerator({ llm: model })
+            })
+        } else if (responseSynthesizerObj.type === 'SimpleResponseBuilder') {
+            const responseSynthesizer = new ResponseSynthesizer({
+                responseBuilder: new SimpleResponseBuilder(serviceContext),
+                serviceContext
+            })
+            queryEngine = SubQuestionQueryEngine.fromDefaults({
+                responseSynthesizer,
+                serviceContext,
+                queryEngineTools,
+                questionGen: new LLMQuestionGenerator({ llm: model })
+            })
+        }
+    }
+
+    return queryEngine
+}
+
 module.exports = { nodeClass: SubQuestionQueryEngine_LlamaIndex }
diff --git a/packages/components/nodes/tools/QueryEngineTool/QueryEngineTool.ts b/packages/components/nodes/tools/QueryEngineTool/QueryEngineTool.ts
@@ -1,5 +1,5 @@
 import { INode, INodeData, INodeParams } from '../../../src/Interface'
-import { VectorStoreIndex } from 'llamaindex'
+import { BaseQueryEngine } from 'llamaindex'
 
 class QueryEngine_Tools implements INode {
     label: string
@@ -16,7 +16,7 @@ class QueryEngine_Tools implements INode {
     constructor() {
         this.label = 'QueryEngine Tool'
         this.name = 'queryEngineToolLlamaIndex'
-        this.version = 1.0
+        this.version = 2.0
         this.type = 'QueryEngineTool'
         this.icon = 'queryEngineTool.svg'
         this.category = 'Tools'
@@ -25,9 +25,9 @@ class QueryEngine_Tools implements INode {
         this.baseClasses = [this.type]
         this.inputs = [
             {
-                label: 'Vector Store Index',
-                name: 'vectorStoreIndex',
-                type: 'VectorStoreIndex'
+                label: 'Base QueryEngine',
+                name: 'baseQueryEngine',
+                type: 'BaseQueryEngine'
             },
             {
                 label: 'Tool Name',
@@ -45,20 +45,15 @@ class QueryEngine_Tools implements INode {
     }
 
     async init(nodeData: INodeData): Promise<any> {
-        const vectorStoreIndex = nodeData.inputs?.vectorStoreIndex as VectorStoreIndex
+        const baseQueryEngine = nodeData.inputs?.baseQueryEngine as BaseQueryEngine
         const toolName = nodeData.inputs?.toolName as string
         const toolDesc = nodeData.inputs?.toolDesc as string
         const queryEngineTool = {
-            queryEngine: vectorStoreIndex.asQueryEngine({
-                preFilters: {
-                    ...(vectorStoreIndex as any).metadatafilter
-                }
-            }),
+            queryEngine: baseQueryEngine,
             metadata: {
                 name: toolName,
                 description: toolDesc
-            },
-            vectorStoreIndex
+            }
         }
 
         return queryEngineTool

diff --git a/packages/server/marketplaces/chatflows/Query Engine.json b/packages/server/marketplaces/chatflows/Query Engine.json
@@ -16,10 +16,10 @@
             "data": {
                 "id": "queryEngine_0",
                 "label": "Query Engine",
-                "version": 1,
+                "version": 2,
                 "name": "queryEngine",
                 "type": "QueryEngine",
-                "baseClasses": ["QueryEngine"],
+                "baseClasses": ["QueryEngine", "BaseQueryEngine"],
                 "tags": ["LlamaIndex"],
                 "category": "Engine",
                 "description": "Simple query engine built to answer question over your data, without memory",
@@ -55,10 +55,10 @@
                 },
                 "outputAnchors": [
                     {
-                        "id": "queryEngine_0-output-queryEngine-QueryEngine",
+                        "id": "queryEngine_0-output-queryEngine-QueryEngine|BaseQueryEngine",
                         "name": "queryEngine",
                         "label": "QueryEngine",
-                        "type": "QueryEngine"
+                        "type": "QueryEngine | BaseQueryEngine"
                     }
                 ],
                 "outputs": {},