from smolagents import (InferenceClientModel, CodeAgent, ToolCallingAgent, 
                        DuckDuckGoSearchTool, VisitWebpageTool, FinalAnswerTool,
                        WikipediaSearchTool, PythonInterpreterTool,
                        )

model_id = "Qwen/Qwen2.5-72B-Instruct"
model = InferenceClientModel(model_id)

web_agent = ToolCallingAgent(
            tools=[
                DuckDuckGoSearchTool(),
                VisitWebpageTool(),
                WikipediaSearchTool()
            ], 
            model=model,
            name="search_agent",
            description="Runs web searches for you. Give it your query as an argument.",
        )
python_agent = CodeAgent(
            tools=[
                PythonInterpreterTool()
            ],
            model=model,
            name='python_agent',
            description='Use additional_authorized_imports for you. You need to do actions and help to answer the questions with python code',
            additional_authorized_imports=[
                "json",
                "pandas",
                "numpy",
                "requests",
                "time",
                "datetime",
            ],
            planning_interval=5,
            verbosity_level=2,
            max_steps=15,
        )

class BasicAgent:
    """An agent who is able to answer questions."""
    def __init__(self):

        # Instantiate Agent
        self.agent = CodeAgent(tools=[
            FinalAnswerTool()
        ],
            model=model,
            managed_agents=[
                web_agent,
                python_agent
            ],
        additional_authorized_imports=["requests"],
        )
        
        print("BasicAgent initialized.")
        
    def __call__(self, question: str) -> str:
        print(f"Agent received question: {question}...")
        answer = self.agent.run(question)
        print(f"Agent returning answer: {answer}")
        
        return answer