Update agent.py

agent.py

@@ -1,26 +1,23 @@
-from smolagents import CodeAgent, InferenceClientModel
-from tools import all_tools
-
-# define GAIA specific system prompt
-GAIA_SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question. 
-Report your thoughts, and finish your answer with your final response.
-Your final answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
-If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
-If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
-If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
-
-Return ONLY the final answer as the output of your last step."""
+import yaml
+from smolagents import CodeAgent, InferenceClientModel, final_answer
+from tools import search_tool, visit_webpage, handle_file
 
 # model
 model = InferenceClientModel(
     model_id="Qwen/Qwen2.5-Coder-1.5B-Instruct"
 )
+# load the prompt template from yaml file
+with open("prompts.yaml", 'r') as stream:
+    prompt_templates = yaml.safe_load(stream)
 
 # Initialize agent
 agent = CodeAgent(
     model=model,
-    tools=all_tools,
-    additional_authorized_imports=["pandas", "numpy"],
+    tools=[search_tool, visit_webpage, handle_file],
+    max_steps-12,
+    verbosity_level=1,
+    additional_authorized_imports=["pandas", "numpy", "re", "math"],
+    prompt_templates=prompt_templates
 )
 
 # run function
@@ -28,6 +25,5 @@ def run_agent(question: str) -> str:
     """
     This function is used to run the agent. Function is called in app.py
     """
-    full_query = f"{GAIA_SYSTEM_PROMPT}\n\nQuestion:{question}"
     result = agent.run(full_query)
     return str(result).strip()