Update core/chatbot/retrieval_chatbot.py

core/chatbot/retrieval_chatbot.py

@@ -34,7 +34,7 @@ class RetrievalChatbot(BaseChatbot):
         self.summarizer = summarizer if summarizer \
                                      else SimpleRefiner(model=GPT4Model(), sys_prompt=SummaryPrompt.content)
     
-    async def response(self, message: str, image_paths=None, return_logs=False) -> str:
+    async def response(self, message: str, image_paths=None) -> str:
         time1 = time.time()
         print("Query: {message}".format(message=message))
         question = self.decomposer.refine(message, None, image_paths)
@@ -51,10 +51,20 @@ class RetrievalChatbot(BaseChatbot):
             task = asyncio.create_task(self.subquestion_answerer(sub_question, image_paths))
             tasks.append(task)
         results = await asyncio.gather(*tasks)
-        references = "".join(results)
+
+        references = ""
+        all_titles = set([])
+        for result in results:
+            references += result["answer"]
+            for t in result["titles"]:
+                all_titles.add(t)
+        logs = references
+
         time3 = time.time()
         print("Sub references are ",references)
-        refs = self.retriever.retrieve(message)
+        refs, titles = self.retriever.retrieve(message)
+        for t in titles:
+            all_titles.add(t)
         for ref in refs:
             references += "Related research for the user query: {ref}\n".format(ref=ref)
 
@@ -80,16 +90,17 @@ class RetrievalChatbot(BaseChatbot):
         print(f"Decompose: {time2-time1}")
         print(f"Answer Subquestions: {time3-time2}")
         print(f"Summarize: {time4-time3}")
-        
-        if return_logs:
-            return answer, references
-        else:
-            return answer
+
+        return {
+            "answer": answer,
+            "titles": all_titles,
+            "logs": logs
+        }
         
     async def subquestion_answerer(self, sub_question: str, image_paths=None, return_logs=False) -> str:
         sub_retrieve_reference=""
         time_s = time.time()
-        sub_retrieve = self.retriever.retrieve(sub_question)
+        sub_retrieve, titles = self.retriever.retrieve(sub_question)
         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)
@@ -99,4 +110,7 @@ class RetrievalChatbot(BaseChatbot):
         time_e = time.time()
         print(f"Time: {time_e-time_s}")
         print(f"Subanswer: {sub_answer}")
-        return "Subquestion: {sub_question}\nSubanswer: {sub_answer}\n\n\n".format(sub_question=sub_question, sub_answer=sub_answer)
+        return {
+            "answer": "Subquestion: {sub_question}\nSubanswer: {sub_answer}\n\n\n".format(sub_question=sub_question, sub_answer=sub_answer), 
+            "titles": titles
+        }