Update agent.py

agent.py

@@ -1,6 +1,6 @@
 
 
-#from langchain_core.output_parsers import PydanticOutputParser
+from langchain_core.output_parsers import PydanticOutputParser
 from typing import Callable, Dict, List, Any
 import time
 import json
@@ -42,8 +42,8 @@ def format_gaia_response(model_type, last_observation, question_out):
     get_llm_response = select_model(model_type)
 
     # Process Gaia 
-    # with open(base_dir+"/system_prompt_final.txt", "r") as f:
-    #     final_sys_prompt = f.read()
+    with open(base_dir+"/system_prompt_final.txt", "r") as f:
+        final_sys_prompt = f.read()
 
     gaia_prompt = (
         f"{final_sys_prompt}\n\n"
@@ -52,7 +52,7 @@ def format_gaia_response(model_type, last_observation, question_out):
         "Please review user questions and the last obervation and respond with the correct answer, in the correct format. No extra text, just the answer."
     )
 
-    final_answer_out = get_llm_response(model_type, gaia_prompt, reasoning_format = 'hidden')
+    final_answer_out = get_llm_response(final_sys_prompt, gaia_prompt, reasoning_format = 'hidden')
 
     return final_answer_out
 
@@ -65,11 +65,11 @@ class ImprovedAgent:
 
 
         # Load system prompts from .txt files
-        #self.system_prompt_plan = self.load_prompt(base_dir+"/system_prompt_planning.txt")
-        #self.system_prompt_thought = self.load_prompt(base_dir+"/system_prompt_thought.txt")
-        #self.system_prompt_action = self.load_prompt(base_dir+"/system_prompt_action.txt")
-        #self.system_prompt_observe = self.load_prompt(base_dir+"/system_prompt_observe.txt")
-        self.system_prompt_main = self.load_prompt(base_dir+"/system_prompt_main.txt")
+        self.system_prompt_plan = self.load_prompt(base_dir+"/system_prompt_planning.txt")
+        self.system_prompt_thought = self.load_prompt(base_dir+"/system_prompt_thought.txt")
+        self.system_prompt_action = self.load_prompt(base_dir+"/system_prompt_action.txt")
+        self.system_prompt_observe = self.load_prompt(base_dir+"/system_prompt_observe.txt")
+
 
     def load_prompt(self, filepath: str) -> str:
         with open(filepath, "r") as f:
@@ -122,7 +122,7 @@ History: {json.dumps(self.history, indent=2)}
         planning_input = f"User Query: {query}"
         print("-----Stage Plan-----")
         
-        plan_response = self.get_llm_response(self.system_prompt_main, planning_input)
+        plan_response = self.get_llm_response(self.system_prompt_plan, planning_input)
         print("-----Plan Text-----")
         print(plan_response)
         print("-------------------")
@@ -140,7 +140,7 @@ History: {json.dumps(self.history, indent=2)}
             # Step 2: Thought Agent
             print("-----Stage Thought-----")
             
-            thought_response = self.get_llm_response(self.system_prompt_main, current_input)
+            thought_response = self.get_llm_response(self.system_prompt_thought, current_input)
             print(thought_response)
             parsed_thought = self.parse_json_response(thought_response)
             print("-----Thought Parsed-----")
@@ -217,7 +217,7 @@ History: {json.dumps(self.history, indent=2)}
                                     Tool Output: {observation}
                                     """         
             print("-----Stage Observe-----")
-            observation_response_text = self.get_llm_response(self.self.system_prompt_main, observation_input)
+            observation_response_text = self.get_llm_response(self.system_prompt_observe, observation_input)
 
             print("-----Observation Parsed-----")
             parsed_observation = self.parse_json_response(observation_response_text)