Refactored from agent to multi-agent

agent.py

@@ -19,21 +19,54 @@ BASE_PYTHON_TOOLS["io"] = io
 BASE_PYTHON_TOOLS["contextlib"] = contextlib
 BASE_PYTHON_TOOLS["exec"] = exec
 
-
-class BasicAgent:
+class ResearchAgent:
     def __init__(self, settings: Settings):
         self.agent = CodeAgent(
+            name="researcher",
+            description="Searches the web, works with files, and answers questions for you. Give it your query as an argument.",
             add_base_tools=False,
             tools=[GoogleSearchTool("serper"),
                    VisitWebpageTool(max_output_length=100000),
-                   FinalAnswerTool(),
-                   GetTaskFileTool(settings),
                    VideoUnderstandingTool(settings, GoogleModelID.GEMINI_2_0_FLASH),
-                   AudioUnderstandingTool(settings, GoogleModelID.GEMINI_2_0_FLASH),
-                   ChessPiecePlacementTool(),
+                   AudioUnderstandingTool(settings, GoogleModelID.GEMINI_2_0_FLASH)
+                   ],
+            additional_authorized_imports=[
+                "unicodedata",
+                "stat",
+                "datetime",
+                "random",
+                "pandas",
+                "itertools",
+                "math",
+                "statistics",
+                "queue",
+                "time",
+                "collections",
+                "re",
+                "os"
+            ],
+            max_steps=10,
+            verbosity_level=1,
+            model=LiteLLMModel(
+                model_id=OpenRouterModelID.GPT_O4_MINI,
+                # model_id=OpenRouterModelID.GPT_4_1_MINI,
+                # model_id=OpenRouterModelID.GROK_3_BETA,
+                # model_id=OpenRouterModelID.GROK_3_MINI_BETA,
+                api_key = settings.openrouter_api_key.get_secret_value(),
+                temperature=0.0, timeout=180
+            )
+        )
+
+class ChessAgent:
+    def __init__(self, settings: Settings):
+        self.agent = CodeAgent(
+            name="chess_player",
+            description="Makes a chess move. Give it a query including board image filepath and player turn (black or white).",
+            add_base_tools=False,
+            tools=[ChessPiecePlacementTool(),
                    ChessGameFenTool(settings, OpenRouterModelID.GPT_O4_MINI),
                    BestChessMoveTool(settings),
-                   ConvertChessMoveTool(settings, OpenRouterModelID.QWEN_3_14B_FREE)
+                   ConvertChessMoveTool(settings, OpenRouterModelID.QWEN_3_14B_FREE),
                    ],
             additional_authorized_imports=[
                 "unicodedata",
@@ -53,13 +86,28 @@ class BasicAgent:
             max_steps=10,
             verbosity_level=1,
             model=LiteLLMModel(
-                # model_id=OpenRouterModelID.GPT_O4_MINI,
-                model_id=OpenRouterModelID.GPT_4_1_MINI,
+                model_id=OpenRouterModelID.GPT_O4_MINI,
+                api_key = settings.openrouter_api_key.get_secret_value(),
+                temperature=0.0, timeout=180
+            )
+        )
+
+class ManagerAgent:
+    def __init__(self, settings: Settings):
+        self.researcher = ResearchAgent(settings).agent
+        self.chess_player = ChessAgent(settings).agent
+        self.agent = CodeAgent(
+            tools=[GetTaskFileTool(settings), FinalAnswerTool()],
+            model=LiteLLMModel(
+                model_id=OpenRouterModelID.GPT_O4_MINI,
+                # model_id=OpenRouterModelID.GPT_4_1_MINI,
                 # model_id=OpenRouterModelID.GROK_3_BETA,
                 # model_id=OpenRouterModelID.GROK_3_MINI_BETA,
+                # model_id=OpenRouterModelID.QWEN_3_14B_FREE,
                 api_key = settings.openrouter_api_key.get_secret_value(),
                 temperature=0.0, timeout=180
-            )
+            ),
+            managed_agents=[self.researcher, self.chess_player],
         )
         # print("BasicAgent initialized.")
     def __call__(self, question: str) -> str:

runner.py

@@ -1,6 +1,6 @@
 from settings import Settings
 from models import Question, QuestionAnswerPair
-from agent import BasicAgent
+from agent import ManagerAgent
 import pandas as pd
 import logging
 import json
@@ -9,7 +9,6 @@ import nest_asyncio
 nest_asyncio.apply()
 logger = logging.getLogger(__name__)
 
-
 class Runner():
     def __init__(self, settings: Settings):
         self.settings = settings
@@ -28,8 +27,11 @@ class Runner():
             f"{item.question} "
             "Think hard to answer. Parse all statements in the question to make a plan. "
             "Your final answer should be a number or as few words as possible. "
-            "If needed, use a comma separated list of numbers and/or strings. Critically "
-            f"review your answer before making it the final answer. task_id: {task_id}."
+            "Only use abbreviations when the question calls for abbreviations. "
+            "If needed, use a comma separated list of values; the comma is always followed by a space"
+            f"Critically review your answer before making it the final answer. "
+            f"Double check the answer to make sure it meets all format requirements stated in the question. "
+            f"task_id: {task_id}."
         )
         if file_name:
             question_text = f"{question_text} file_name: {file_name} (use tools to fetch the file)"
@@ -40,7 +42,7 @@ class Runner():
         task_id = item.task_id
         question_text = self._enrich_question_text(item)
         try:
-            answer = await asyncio.to_thread(BasicAgent(self.settings), question_text)
+            answer = await asyncio.to_thread(ManagerAgent(self.settings), question_text)
         except Exception as e:
             logger.error(f"Error running agent on task {task_id}: {e}")
             answer = f"AGENT ERROR: {e}"

tools.py

@@ -21,10 +21,10 @@ class BaseCustomTool(Tool):
         
 class GetTaskFileTool(BaseCustomTool):
     name = "get_task_file_tool"
-    description = """Download the file_name associated with a given task_id. Get absolute file path"""
+    description = """If a file_name is provided, download file associated with a given task_id. Get absolute file path"""
     inputs = {
-        "task_id": {"type": "string", "description": "Task ID"},
-        "file_name": {"type": "string", "description": "File name"},
+        "task_id": {"type": "string", "description": "Task ID (required)"},
+        "file_name": {"type": "string", "description": "File name (required)"},
     }
     output_type = "string"