From ba554bceba91f95de7eb6d858674e393c4c57a23 Mon Sep 17 00:00:00 2001
From: didi <2020201387@ruc.edu.cn>
Date: Thu, 11 Apr 2024 00:03:23 +0800
Subject: [PATCH] update math resolver & prompt

---
 math_ai/codebase/gate_controller.py | 13 +++++++++++--
 math_ai/codebase/math_resovler.py   | 18 ++++++------------
 math_ai/codebase/prompt.py          |  4 +---
 math_ai/codebase/strategies.py      | 16 ++++++++++++++--
 4 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/math_ai/codebase/gate_controller.py b/math_ai/codebase/gate_controller.py
index a0c12f4563..99688e2217 100644
--- a/math_ai/codebase/gate_controller.py
+++ b/math_ai/codebase/gate_controller.py
@@ -13,18 +13,27 @@ def __init__(self):
         self.llm = OpenAILLM()
         self.role = "<TODO Here is Gate Controller's system prompt>"
         self.llm.set_role(self.role)
+
     def run(self, problem: Dict) -> Dict:
         """
+        problem: 
+        {
+            "desc":"<QUESTION DESC>",
+            "type":"<select from human design>",
+        }
+
         Gate Controller choose human design strategy here
         and return strategy in dict
         you can add anything else such as problem's possible attention in dict.
         for example:
+        
+        return 
         {
-            "strategy": "<content>",
+            "strategy": "策略名称",
             "attention": "<"Attention points identified during the determination of problem types.">"
         }
         """
 
-        return {"strategy":"<content>"}
+        return {"strategy":"知识类型"}
 
 
diff --git a/math_ai/codebase/math_resovler.py b/math_ai/codebase/math_resovler.py
index caff11552c..5931545b38 100644
--- a/math_ai/codebase/math_resovler.py
+++ b/math_ai/codebase/math_resovler.py
@@ -9,12 +9,10 @@
 from metagpt.roles.di.data_interpreter import DataInterpreter
 from math_ai.codebase.engine.llm import OpenAILLM
 from math_ai.codebase.prompt import zero_shot_planner, resolver_planner, inference_prompt, di_prompt, result_validate_prompt, inference_final_prompt, logic_validate_prompt
-
+from math_ai.codebase.strategies import get_strategy_desc
 # TODO add different phase in codebase.phase
 
-async def main(requirement: str = ""):
-    di = DataInterpreter()
-    await di.run(requirement)
+
 
 class MathResolver:
     def __init__(self):
@@ -37,8 +35,8 @@ async def single_run(self, problem: Dict, types: Dict) -> Dict:
         Finally, math resolver need to return the solution without refine.
         """
         strategy_name = types["strategy"]
-        strategy = ""
-
+        strategy = get_strategy_desc(strategy_name)
+        
         # 1. 直接要求他解决数学问题，思考这个过程。 zero shot 让他先去对这个题目给出一个计划。
         # 2. 得到这个过程之后，让他结合我们的strategy 跟 Prompt，重新构建phase
         # 3. 每一个Phase的Prompt如何去写 
@@ -49,21 +47,16 @@ async def single_run(self, problem: Dict, types: Dict) -> Dict:
         current_trajectory = ""
         for phase in resolver_plan:
             if phase["plan"]["phase"] == "inference":
-                # TODO 这里的Prompt 需要修改
                 current_trajectory += self.inference(problem, current_trajectory, subgoal=phase["plan"]["desc"])
             elif phase["plan"]["phase"] == "di":
                 current_trajectory += self.di_run(problem, current_trajectory, subgoal=phase["plan"]["desc"])
             elif phase["plan"]["phase"] == "logic_validate":
-                # TODO 如果Validate 失败，是否需要重新进行plan
                 current_trajectory += self.logic_validate(problem, current_trajectory, subgoal=phase["plan"]["desc"])
         
-        # TODO 在这里result validate
         if self.result_validate(problem, current_trajectory):
             pass
         else:
             current_trajectory += self.inference_final(problem, current_trajectory)
-
-            pass
         return {"current_trajectory": current_trajectory}
     
     async def multi_run(self):
@@ -89,4 +82,5 @@ def result_validate(self, problem, current_trajectory):
     
     def inference_final(self, problem, current_trajectory):
         final_result = self.llm.llm_response(prompt=inference_final_prompt.format(problem=problem, trajectoty=current_trajectory))
-        return final_result
\ No newline at end of file
+        return final_result
+    
diff --git a/math_ai/codebase/prompt.py b/math_ai/codebase/prompt.py
index c0a971788d..94662eccfc 100644
--- a/math_ai/codebase/prompt.py
+++ b/math_ai/codebase/prompt.py
@@ -68,10 +68,8 @@
 """
 
 di_prompt = """
-
 在解决{problem}的过程中，你被指派解决这一子问题{subgoal}。
-
-
+请你注意，在规划的过程中，需要专注于子问题，大问题只是用来给你提供更多的背景信息。
 """
 
 
diff --git a/math_ai/codebase/strategies.py b/math_ai/codebase/strategies.py
index 7f25fb22d7..e7a19d1af7 100644
--- a/math_ai/codebase/strategies.py
+++ b/math_ai/codebase/strategies.py
@@ -26,6 +26,18 @@
 
 
 """
-
-
+def get_strategy_desc(name):
+    if name == "ALGEBRA_NUMBER_THEORY":
+        return ALGEBRA_NUMBER_THEORY
+    elif name == "GEOMOTRY_TOPOLOGY":
+        return GEOMOTRY_TOPOLOGY
+    elif name == "ANALYSYS_EQUATION":
+        return ANALYSYS_EQUATION
+    elif name == "COMBINATION_PROBABILITY":
+        return COMBINATION_PROBABILITY
+    elif name == "COMPUTATION":
+        return COMPUTATION
+    else:
+        return "No Strategy"
+    
 print(ANALYSYS_EQUATION)
\ No newline at end of file