agent.py

import json
from pathlib import Path
from smolagents import Agent
from smolagents.core.task import Task
from mistral_hf_wrapper import HuggingFaceModel

# Import tools
from tools import (
    chess_tools,
    classifier_tool,
    content_retriever_tool,
    get_attachments_tool,
    google_search_tools,
    speech_recognition_tool,
    youtube_video_tool,
)

# Load GAIA task list from metadata.jsonl
def load_tasks(path: str = "metadata.jsonl") -> list[Task]:
    with open(path, "r", encoding="utf-8") as f:
        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"
)

# List of tool instances
tools = [
    chess_tools.ImageToChessBoardFENTool(),
    classifier_tool.ClassifierTool(),
    content_retriever_tool.ContentRetrieverTool(),
    get_attachments_tool.GetAttachmentTool(),
    google_search_tools.GoogleSearchTool(),
    google_search_tools.GoogleSiteSearchTool(),
    speech_recognition_tool.SpeechRecognitionTool(),
    youtube_video_tool.YouTubeVideoTool(),
]

# Build the agent
agent = Agent(
    model=model,
    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)