From 3d670fec817e8d7b0d4b6de5ee262ffc5b64efa8 Mon Sep 17 00:00:00 2001
From: Henry <hzj94@hotmail.com>
Date: Fri, 26 Jan 2024 00:22:24 +0000
Subject: [PATCH] update query engine

---
 .../nodes/engine/QueryEngine/QueryEngine.ts          | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/packages/components/nodes/engine/QueryEngine/QueryEngine.ts b/packages/components/nodes/engine/QueryEngine/QueryEngine.ts
index 14774703..bd6e040d 100644
--- a/packages/components/nodes/engine/QueryEngine/QueryEngine.ts
+++ b/packages/components/nodes/engine/QueryEngine/QueryEngine.ts
@@ -68,13 +68,15 @@ class QueryEngine_LlamaIndex implements INode {
         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
                 })
-                return new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
+                queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
             } else if (responseSynthesizerObj.type === 'CompactAndRefine') {
                 const responseSynthesizer = new ResponseSynthesizer({
                     responseBuilder: new CompactAndRefine(
@@ -84,7 +86,7 @@ class QueryEngine_LlamaIndex implements INode {
                     ),
                     serviceContext: vectorStoreRetriever.serviceContext
                 })
-                return new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
+                queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
             } else if (responseSynthesizerObj.type === 'Refine') {
                 const responseSynthesizer = new ResponseSynthesizer({
                     responseBuilder: new Refine(
@@ -94,18 +96,16 @@ class QueryEngine_LlamaIndex implements INode {
                     ),
                     serviceContext: vectorStoreRetriever.serviceContext
                 })
-                return new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
+                queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
             } else if (responseSynthesizerObj.type === 'SimpleResponseBuilder') {
                 const responseSynthesizer = new ResponseSynthesizer({
                     responseBuilder: new SimpleResponseBuilder(vectorStoreRetriever.serviceContext),
                     serviceContext: vectorStoreRetriever.serviceContext
                 })
-                return new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
+                queryEngine = new RetrieverQueryEngine(vectorStoreRetriever, responseSynthesizer)
             }
         }
 
-        const queryEngine = new RetrieverQueryEngine(vectorStoreRetriever)
-
         let text = ''
         let sourceDocuments: ICommonObject[] = []
         let sourceNodes: BaseNode<Metadata>[] = []