Add additional tools; tune agent system prompts

agent.py

@@ -2,18 +2,22 @@ from smolagents import (
     ToolCallingAgent,
     CodeAgent,
     DuckDuckGoSearchTool,
+    PythonInterpreterTool,
     FinalAnswerTool,
     VisitWebpageTool,
     PromptTemplates,
     PlanningPromptTemplate,
     ManagedAgentPromptTemplate,
     FinalAnswerPromptTemplate,
+    WikipediaSearchTool,
 )
 
 from tool import addition_tool
 from src.config.model import openai_model
+import re
 
 
+gaia_prompt = "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. 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 then reply with digit and not in words, e.g: if Five, then reply with 5, if Three, then reply with 3, 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."
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 class BaseAgent:
@@ -22,14 +26,18 @@ class BaseAgent:
         self.agent = ToolCallingAgent(
             model=openai_model,
             tools=[
-                # addition_tool,
+                addition_tool,
+                WikipediaSearchTool(),
+                PythonInterpreterTool(),
                 DuckDuckGoSearchTool(),
-                # VisitWebpageTool(),
-                # FinalAnswerTool(),
+                VisitWebpageTool(),
+                FinalAnswerTool(),
             ],
+            add_base_tools=True,
+            max_steps=4,
             verbosity_level=2,
             prompt_templates=PromptTemplates(
-                system_prompt="You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. 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.",
+                system_prompt=gaia_prompt,
                 planning=PlanningPromptTemplate(
                     initial_plan="",
                     update_plan_pre_messages="",
@@ -40,7 +48,7 @@ class BaseAgent:
                     report="",
                 ),
                 final_answer=FinalAnswerPromptTemplate(
-                    pre_messages="", post_messages=""
+                    pre_messages=gaia_prompt, post_messages=gaia_prompt+"The final answer should be as few words or numbers only. Exclude all the previuos context and only return the final answer. In case words separated by commas, use the following format: 'word1, word2, word3' without any additional text or explanation.",
                 ),
             ),
         )
@@ -48,5 +56,7 @@ class BaseAgent:
     def run(self, question: str, additional_args: dict = None) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         response = self.agent.run(question, additional_args=additional_args)
+        response = re.sub(r"(?s)^.*?FINAL ANSWER:", "FINAL ANSWER:", response)
+        response = response.replace("FINAL ANSWER:", "").strip()
         print(f"Agent returning response: {response}")
-        return response.replace("FINAL ANSWER:", "").strip()
+        return response

requirements.txt

@@ -2,3 +2,4 @@ gradio
 requests
 smolagents[toolkit,openai,vision,transformers,litellm]
 gradio[oauth]
+wikipedia-api