agent.py

# agent.py
import contextlib
import io
import logging
import os
logger = logging.getLogger(__name__)

from models import GoogleModelID # Import GoogleModelID
from settings import Settings
from smolagents import OpenAIServerModel, CodeAgent, FinalAnswerTool # Changed from LiteLLMModel
from smolagents import DuckDuckGoSearchTool, VisitWebpageTool # Changed from GoogleSearchTool
from smolagents.local_python_executor import BASE_PYTHON_TOOLS
from tools import GetTaskFileTool, VideoUnderstandingTool, AudioUnderstandingTool
from tools import ChessBoardFENTool, BestChessMoveTool, ConvertChessMoveTool, ExcelParsingTool
import json # Added for BASE_PYTHON_TOOLS
import pandas as pd # Added for BASE_PYTHON_TOOLS


# Extend BASE_PYTHON_TOOLS for the PythonInterpreterTool to have access to these
BASE_PYTHON_TOOLS["open"] = open
BASE_PYTHON_TOOLS["os"] = os
BASE_PYTHON_TOOLS["io"] = io
BASE_PYTHON_TOOLS["contextlib"] = contextlib
BASE_PYTHON_TOOLS["exec"] = exec # Note: exec is powerful, use with caution in production
BASE_PYTHON_TOOLS["json"] = json # For parsing JSON if needed by agent
BASE_PYTHON_TOOLS["pd"] = pd # For pandas operations if needed by agent

class ResearchAgent:
    def __init__(self, settings: Settings):
        self.agent = CodeAgent(
            name="researcher",
            description="A specialized agent for web research, video analysis, and audio understanding. Give it your query as an argument. Use 'duckduckgo_search_tool' for web searches, 'visit_webpage_tool' to read web page content, 'video_understanding_tool' for YouTube videos, and 'audio_understanding_tool' for local audio files.",
            add_base_tools=False,
            tools=[
                DuckDuckGoSearchTool(), # Changed from GoogleSearchTool
                VisitWebpageTool(max_output_length=100000),
                VideoUnderstandingTool(settings, GoogleModelID.GEMINI_2_0_FLASH), # Still uses 2.0 Flash for specific multimodal tasks
                AudioUnderstandingTool(settings, GoogleModelID.GEMINI_2_0_FLASH) # Still uses 2.0 Flash for specific multimodal tasks
            ],
            additional_authorized_imports=[
                "unicodedata", "stat", "datetime", "random", "pandas", "itertools",
                "math", "statistics", "queue", "time", "collections", "re", "os",
                "json", "io", "urllib.parse"
            ],
            max_steps=15,
            verbosity_level=2,
            model=OpenAIServerModel( # Changed to OpenAIServerModel
                model_id=GoogleModelID.GEMINI_2_5_FLASH_PREVIEW, # Set to GEMINI_2_5_FLASH_PREVIEW
                api_base="https://generativelanguage.googleapis.com/v1beta/openai/", # Gemini API base
                api_key = settings.gemini_api_key.get_secret_value(), # Use Gemini API key
                temperature=0.1,
                timeout=180
            )
        )
        logger.info("ResearchAgent initialized.")

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=[
                ChessBoardFENTool(),
                BestChessMoveTool(settings),
                ConvertChessMoveTool(settings, GoogleModelID.GEMINI_2_5_FLASH_PREVIEW), # Changed to Gemini Flash Preview
            ],
            additional_authorized_imports=[
                "unicodedata", "stat", "datetime", "random", "pandas", "itertools",
                "math", "statistics", "queue", "time", "collections", "re", "os",
                "json", "urllib.parse"
            ],
            max_steps=10,
            verbosity_level=2,
            model=OpenAIServerModel( # Changed to OpenAIServerModel
                model_id=GoogleModelID.GEMINI_2_5_FLASH_PREVIEW, # Set to GEMINI_2_5_FLASH_PREVIEW
                api_base="https://generativelanguage.googleapis.com/v1beta/openai/", # Gemini API base
                api_key = settings.gemini_api_key.get_secret_value(), # Use Gemini API key
                temperature=0.0,
                timeout=180
            )
        )
        logger.info("ChessAgent initialized.")

class ManagerAgent:
    """
    The main orchestrating agent that routes questions to specialized sub-agents
    or handles them directly with its own tools.
    """
    def __init__(self, settings: Settings):
        self.settings = settings
        self.researcher = ResearchAgent(settings).agent
        self.chess_player = ChessAgent(settings).agent
        
        # Main manager agent
        self.agent = CodeAgent(
            name="manager",
            description=(
                "You are a highly capable AI assistant designed to solve complex GAIA benchmark questions. "
                "Your primary role is to route tasks to the most appropriate specialized agent: "
                "'researcher' for general knowledge, web browsing, video, and audio understanding tasks, "
                "or 'chess_player' for chess-related tasks. "
                "If a task involves downloading a file, use 'get_task_file_tool' first. "
                "If you have the final answer, use 'final_answer_tool'.\n\n"
                "**Available Tools:**\n"
                "- `get_task_file_tool(task_id: str, file_name: str)`: Downloads a file associated with a task.\n"
                "- `final_answer_tool(answer: str)`: Use this when you have the exact final answer.\n\n"
                "**Managed Agents:**\n"
                "- `researcher(query: str)`: Use for questions requiring web search, video analysis, or audio analysis.\n"
                "- `chess_player(query: str)`: Use for questions related to chess positions or moves.\n\n"
                "Think step-by-step. If a task involves a file, use `get_task_file_tool` first to download it, then pass the file path to the appropriate sub-agent or tool."
            ),
            tools=[
                GetTaskFileTool(settings),
                FinalAnswerTool(),
                ExcelParsingTool(settings) # Added ExcelParsingTool to ManagerAgent as it handles file paths
            ],
            model=OpenAIServerModel( # Changed to OpenAIServerModel
                model_id=GoogleModelID.GEMINI_2_5_FLASH_PREVIEW, # Set to GEMINI_2_5_FLASH_PREVIEW
                api_base="https://generativelanguage.googleapis.com/v1beta/openai/", # Gemini API base
                api_key = settings.gemini_api_key.get_secret_value(), # Use Gemini API key
                temperature=0.0,
                timeout=180
            ),
            managed_agents=[self.researcher, self.chess_player],
            verbosity_level=2,
            max_steps=20
        )
        logger.info("ManagerAgent initialized.")

    def __call__(self, question_data: dict) -> str:
        task_id = question_data.get("task_id", "N/A")
        question_text = question_data.get("question", "")
        file_name = question_data.get("file_name", "")

        enriched_question = (
            f"{question_text} "
            f"task_id: {task_id}. "
            f"Your final answer should be a number or as few words as possible. "
            f"Only use abbreviations when the question calls for abbreviations. "
            f"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. "
        )
        if file_name:
            enriched_question = f"{enriched_question} file_name: {file_name} (use get_task_file_tool to fetch this file and then pass its path to the relevant tool/agent, or excel_parsing_tool if it's an Excel file)." # Updated prompt for Excel
        
        logger.info(f"ManagerAgent received question (first 100 chars): {enriched_question[:100]}...")
        
        try:
            final_answer = self.agent.run(enriched_question)
            logger.info(f"ManagerAgent returning final answer: {final_answer}")
            return final_answer
        except Exception as e:
            logger.error(f"Error running ManagerAgent on task {task_id}: {e}")
            return f"AGENT ERROR: {e}"