Update core/chatbot/retrieval_chatbot.py

core/chatbot/retrieval_chatbot.py

@@ -4,6 +4,7 @@ from ..retriever import BaseRetriever, ChromaRetriever, FaissRetriever
 from ..refiner import BaseRefiner, SimpleRefiner
 from models import BaseModel, GPT4Model
 from prompts import DecomposePrompt, QAPrompt, SummaryPrompt, ReferencePrompt
+from utils import convert_str_to_list
 import ast
 from utils.image_encoder import encode_image
 import asyncio
@@ -39,12 +40,15 @@ class RetrievalChatbot(BaseChatbot):
         print("Query: {message}".format(message=message))
         question = self.decomposer.refine(message, None, image_paths)
         print(question)
-        question = question.replace('"', "'").replace("\n", "").replace("', '", "','").lstrip("['").rstrip("']")
-        sub_questions = question.split("','")
-        print("Decomposed your query into subquestions: {sub_questions}".format(sub_questions=sub_questions))
+        # question = question.replace('"', "'").replace("\n", "").replace("', '", "','").lstrip("['").rstrip("']")
+        # sub_questions = question.split("','")
+        # print("Decomposed your query into subquestions: {sub_questions}".format(sub_questions=sub_questions))
+        sub_questions_str = self.decomposer.refine(message, None, image_paths)
+        sub_questions_list = convert_str_to_list(sub_questions_str)
+        print("Decomposed your query into subquestions: {sub_questions}".format(sub_questions=sub_questions_list))
         tasks = []
         time2 = time.time()
-        for sub_question in sub_questions:
+        for sub_question in sub_questions_list:
             # print("="*20)
             # print(f"Subquestion: {sub_question}")            
             # print(f"Retrieving pdf papers for references...\n")
@@ -100,7 +104,16 @@ class RetrievalChatbot(BaseChatbot):
     async def subquestion_answerer(self, sub_question: str, image_paths=None, return_logs=False) -> str:
         sub_retrieve_reference=""
         time_s = time.time()
-        sub_retrieve, titles = self.retriever.retrieve(sub_question)
+        # sub_retrieve, titles = self.retriever.retrieve(sub_question)
+        if len(sub_question) >= 2:
+            sub_retrieve_english, titles_english = self.retriever.retrieve(sub_question[0])
+            sub_retrieve_chinese, titles_chinese = self.retriever.retrieve(sub_question[1])
+        else:
+            sub_retrieve_english, titles_english = self.retriever.retrieve(sub_question[0])
+            sub_retrieve_chinese, titles_chinese = [], []
+        # print("英文检索结果: ", sub_retrieve_english)
+        # print("中文检索结果: ", sub_retrieve_chinese)
+        sub_retrieve=sub_retrieve_english+sub_retrieve_chinese
         for ref in sub_retrieve:
             sub_retrieve_reference += "Related research: {ref}\n".format(ref=ref)
         sub_answerer_context = "Sub Question References: {sub_retrieve_reference}\nQuestion: {question}\n".format(sub_retrieve_reference=sub_retrieve_reference, question=sub_question)
@@ -112,5 +125,5 @@ class RetrievalChatbot(BaseChatbot):
         print(f"Subanswer: {sub_answer}")
         return {
             "answer": "Subquestion: {sub_question}\nSubanswer: {sub_answer}\n\n\n".format(sub_question=sub_question, sub_answer=sub_answer), 
-            "titles": titles
+            "titles": titles_english + titles_chinese
         }