feat: use retriever instead of query engine in data source tool

l3vels · Oct 31, 2023 · 0d85256 · 0d85256
1 parent 408702a
commit 0d85256
Showing 1 changed file with 14 additions and 9 deletions.
diff --git a/apps/server/datasources/file/file_retriever.py b/apps/server/datasources/file/file_retriever.py
@@ -211,14 +211,19 @@ def query(self, query_str):
             llm=llm, chunk_size=self.chunk_size
         )
 
-        query_engine = self.index.as_query_engine(
-            response_mode=self.response_mode,
-            service_context=service_context,
-            similarity_top_k=self.similarity_top_k
-            if self.index_type == IndexType.VECTOR_STORE.value
-            else None,
-            verbose=True,
+        retriever = self.index.as_retriever(
+            service_context=service_context, verbose=True
         )
 
-        result = query_engine.query(query_str)
-        return result.response
+        # query_engine = self.index.as_query_engine(
+        #     response_mode=self.response_mode,
+        #     service_context=service_context,
+        #     similarity_top_k=self.similarity_top_k
+        #     if self.index_type == IndexType.VECTOR_STORE.value
+        #     else None,
+        #     verbose=True,
+        # )
+
+        nodes = retriever.retrieve(query_str)
+        content = "\n".join([node.text for node in nodes])
+        return content