CoderLakshman/Final_Assignment_Template

1import contextlib2import io3import logging4import os5logger = logging.getLogger(__name__)6from models import GoogleModelID, OpenRouterModelID7from settings import Settings8from smolagents import LiteLLMModel, CodeAgent9from smolagents import GoogleSearchTool, VisitWebpageTool, FinalAnswerTool10from smolagents.local_python_executor import BASE_PYTHON_TOOLS11from tools import GetTaskFileTool, VideoUnderstandingTool, AudioUnderstandingTool12from tools import ChessBoardFENTool, BestChessMoveTool, ConvertChessMoveTool13 14 15# Base tools may use these to process files16BASE_PYTHON_TOOLS["open"] = open17BASE_PYTHON_TOOLS["os"] = os18BASE_PYTHON_TOOLS["io"] = io19BASE_PYTHON_TOOLS["contextlib"] = contextlib20BASE_PYTHON_TOOLS["exec"] = exec21 22class ResearchAgent:23    def __init__(self, settings: Settings):24        self.agent = CodeAgent(25            name="researcher",26            description="Searches the web, works with files, and answers questions for you. Give it your query as an argument.",27            add_base_tools=False,28            tools=[GoogleSearchTool("serper"),29                   VisitWebpageTool(max_output_length=100000),30                   VideoUnderstandingTool(settings, GoogleModelID.GEMINI_2_0_FLASH),31                   AudioUnderstandingTool(settings, GoogleModelID.GEMINI_2_0_FLASH)32                   ],33            additional_authorized_imports=[34                "unicodedata",35                "stat",36                "datetime",37                "random",38                "pandas",39                "itertools",40                "math",41                "statistics",42                "queue",43                "time",44                "collections",45                "re",46                "os"47            ],48            max_steps=10,49            verbosity_level=1,50            model=LiteLLMModel(51                model_id="gemini/gemini-2.0-flash",52                api_key=settings.gemini_api_key.get_secret_value(),53                temperature=0.0,54                timeout=180,55                max_tokens=51256            )57        )58 59class ChessAgent:60    def __init__(self, settings: Settings):61        self.agent = CodeAgent(62            name="chess_player",63            description="Makes a chess move. Give it a query including board image filepath and player turn (black or white).",64            add_base_tools=False,65            tools=[ChessBoardFENTool(),66                   BestChessMoveTool(settings),67                   ConvertChessMoveTool(settings, OpenRouterModelID.GPT_O4_MINI),68                   ],69            additional_authorized_imports=[70                "unicodedata",71                "stat",72                "datetime",73                "random",74                "pandas",75                "itertools",76                "math",77                "statistics",78                "queue",79                "time",80                "collections",81                "re",82                "os"83            ],84            max_steps=10,85            verbosity_level=1,86            model=LiteLLMModel(87                model_id="gemini/gemini-2.0-flash",88                api_key=settings.gemini_api_key.get_secret_value(),89                temperature=0.0,90                timeout=180,91                max_tokens=51292            )93        )94 95class ManagerAgent:96    def __init__(self, settings: Settings):97        self.researcher = ResearchAgent(settings).agent98        self.chess_player = ChessAgent(settings).agent99        self.agent = CodeAgent(100            tools=[GetTaskFileTool(settings), FinalAnswerTool()],101            model=LiteLLMModel(102                model_id="gemini/gemini-2.0-flash",103                api_key=settings.gemini_api_key.get_secret_value(),104                temperature=0.0,105                timeout=180,106                max_tokens=512107            ),108            managed_agents=[self.researcher, self.chess_player],109        )110        # print("BasicAgent initialized.")111    def __call__(self, question: str) -> str:112        logger.info(f"Agent received question (first 50 chars): {question[:50]}...")113        final_answer = self.agent.run(question)114        logger.info(f"Agent returning fixed answer: {final_answer}")115        return final_answer