From da9c5aca823a4f6637704d893c23b3300895c3fa Mon Sep 17 00:00:00 2001
From: YujingYang <yujing.yang2@mail.mcgill.ca>
Date: Mon, 7 Aug 2023 11:40:43 -0400
Subject: [PATCH] reformulate query if the query is duplicate

---
 apps/slackbot/task_agent.py | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/apps/slackbot/task_agent.py b/apps/slackbot/task_agent.py
index 29141370..830ac469 100644
--- a/apps/slackbot/task_agent.py
+++ b/apps/slackbot/task_agent.py
@@ -1,6 +1,8 @@
 from typing import List, Optional
 
 from pydantic import ValidationError
+import openai
+from os import environ
 
 from langchain.chains.llm import LLMChain
 from langchain.chat_models.base import BaseChatModel
@@ -98,7 +100,7 @@ def run(self, task: str) -> str:
         
 
 
-        
+        previous_action = ""
         while True:
             # Discontinue if continuous limit is reached
             loop_count = self.loop_count
@@ -150,6 +152,29 @@ def run(self, task: str) -> str:
             action = self.output_parser.parse(assistant_reply)
             print("action:", action)
             tools = {t.name: t for t in self.tools}
+            if action == previous_action:
+                if action.name == "Search" or action.name == "Context Search":
+                    print("Action name: ", action.name, "\nStart reformulating the query")
+                    instruction = (
+                        f"You want to search for useful information to answer the query: {task}."
+                        f"The original query is: {action.args['query']}"
+                        f"Reformulate the query so that it can be used to search for relevant information."
+                        f"Only return one query instead of multiple queries."
+                        f"Reformulated query:\n\n"
+                    )
+                    openai.api_key = environ.get("OPENAI_KEY")
+                    response = openai.Completion.create(
+                                    engine='text-davinci-003',
+                                    prompt= " ".join(str(i) for i in self.previous_message) + "\n" + instruction,
+                                    temperature=0.7,
+                                    max_tokens=1024,
+                                    top_p=1,
+                                    frequency_penalty=0,
+                                    presence_penalty=0
+                                )
+                    reformulated_query = response['choices'][0]['text']
+                    action.args['query'] = reformulated_query
+
             if action.name == "finish":
                 self.loop_count = self.max_iterations
                 result = "Finished task. "
@@ -193,6 +218,7 @@ def run(self, task: str) -> str:
 
             # self.memory.add_documents([Document(page_content=memory_to_add)])
             self.previous_message.append(HumanMessage(content=memory_to_add))
+            previous_action = action
 
     def set_user_input(self, user_input: str):
         result = f"Command UserInput returned: {user_input}"