Update agent.py

agent.py

@@ -1,7 +1,10 @@
 import json
 from pathlib import Path
-from smolagents import Agent, Task
+from smolagents.core.agent import Agent
+from smolagents.core.task import Task
 from mistral_hf_wrapper import HuggingFaceModel
+
+# Import tools
 from tools import (
     chess_tools,
     classifier_tool,
@@ -12,19 +15,18 @@ from tools import (
     youtube_video_tool,
 )
 
+# Load GAIA task list from metadata.jsonl
 def load_tasks(path: str = "metadata.jsonl") -> list[Task]:
-    tasks = []
     with open(path, "r", encoding="utf-8") as f:
-        for line in f:
-            task_dict = json.loads(line.strip())
-            tasks.append(Task(**task_dict))
-    return tasks
+        return [Task(**json.loads(line)) for line in f]
 
+# Configure Mistral Inference Endpoint wrapper
 model = HuggingFaceModel(
     endpoint_url="https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"
 )
 
-toolset = [
+# List of tool instances
+tools = [
     chess_tools.ImageToChessBoardFENTool(),
     classifier_tool.ClassifierTool(),
     content_retriever_tool.ContentRetrieverTool(),
@@ -35,11 +37,13 @@ toolset = [
     youtube_video_tool.YouTubeVideoTool(),
 ]
 
+# Build the agent
 agent = Agent(
     model=model,
-    tools=toolset,
-    system_message="You are a helpful assistant solving GAIA benchmark tasks. Return only the final answer, nothing else."
+    tools=tools,
+    system_message="You are a helpful AI agent solving GAIA benchmark tasks. Only return the answer, no extra explanation."
 )
 
+# Core solve function
 def solve_task(task: Task) -> str:
     return agent.run(task.question, task_id=task.task_id)