agent.py

import os
from smolagents import CodeAgent, LiteLLMModel, DuckDuckGoSearchTool, WikipediaSearchTool
from tools import VisitWebpageTool, DownloadTaskAttachmentTool

class TheAgent:
    """Agent that processes questions using LLM and various tools."""

    def __init__(self):
        # Initialize the agent with model and tools
        self.agent = CodeAgent(
            model=LiteLLMModel(
                model_id="openrouter/meta-llama/llama-4-maverick:free",
                api_key=os.getenv("OPENROUTER_KEY")
            ),
            tools=[
                DuckDuckGoSearchTool(), 
                WikipediaSearchTool(), 
                VisitWebpageTool(), 
                DownloadTaskAttachmentTool()
            ],
            add_base_tools=True,
            additional_authorized_imports=['pandas', 'numpy', 'csv', 'subprocess', 'exec']
        )
        print("Agent initialized successfully.")

    def __call__(self, question: str) -> str:
        """Process a question and return the answer."""
        print(f"Processing question: {question[:50]}..." if len(question) > 50 else f"Processing question: {question}")
        return self.agent.run(question)