agent.py

from smolagents import CodeAgent, LiteLLMModel, GoogleSearchTool, WikipediaSearchTool, Tool, VisitWebpageTool, HfApiModel
import os
import json
import requests
import pandas as pd
from huggingface_hub import InferenceClient
from openai import OpenAI

class GetFileTool(Tool):
    name = "get_file_tool"
    description = "This tool allows to download the file attached to the question"
    output_type = "string"
    inputs = {
        "task_id": {
            "type": "string",
            "description": "The task id of the question",
        },
        "file_name": {
            "type": "string",
            "description": "The name of the file to download"
        }
    }

    def forward(self, task_id: str, file_name: str) -> str:
        # Download the file from the task id
        file = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
        # Save the file with the file name
        with open(file_name, "wb") as f:
            f.write(requests.get(file).content)

        # Return the file name  
        return os.path.abspath(file_name)
    
class LoadXlsxFileTool(Tool):
    name = "load_xlsx_file_tool"
    description = """This tool loads xlsx file into pandas and returns it"""
    inputs = {
        "file_path": {"type": "string", "description": "File path"}
    }
    output_type = "object"

    def forward(self, file_path: str) -> object:
        return pd.read_excel(file_path)

class LoadTextFileTool(Tool):
    name = "load_text_file_tool"
    description = """This tool loads any text file"""
    inputs = {
        "file_path": {"type": "string", "description": "File path"}
    }
    output_type = "string"

    def forward(self, file_path: str) -> object:
        with open(file_path, 'r', encoding='utf-8') as file:
            return file.read() 

class AudioToTextTool(Tool):
    name = "audio_to_text_tool"
    description = """This tool transcribes audio files into text"""
    inputs = {
        "file_path": {"type": "string", "description": "File path"}
    }
    output_type = "string"

    def forward(self, file_path: str) -> str:
        try:
            # Check if file exists
            if not os.path.exists(file_path):
                return f"Error: File {file_path} does not exist"

            # Read the audio file as raw bytes
            with open(file_path, "rb") as f:
                audio_data = f.read()

            # Set up the API URL and headers
            api_url = "https://router.huggingface.co/hf-inference/models/openai/whisper-large-v3-turbo"
            headers = {
                "Authorization": f"Bearer {os.getenv('HF_TOKEN')}",
                "Content-Type": "audio/mpeg"  # Assuming MP3 format
            }

            # Make the API request
            response = requests.post(api_url, headers=headers, data=audio_data)
            response.raise_for_status()  # Raise an exception for bad status codes
            
            # Parse and return the response
            output = response.json()
            return output["text"]

        except Exception as e:
            return f"Error transcribing audio: {str(e)}"
    
class ImageAnalysisTool(Tool):
    name = "image_analysis_tool"
    description = """This tool analyzes images and returns the text and the information in the image"""
    inputs = {
        "task_id": {"type": "string", "description": "The task id of the question"},
    }
    output_type = "string"

    def forward(self, task_id: str) -> str:
        client = InferenceClient(
            provider="nebius",
            api_key=os.getenv("HF_TOKEN"),
        )
            
        completion = client.chat.completions.create(
            model="google/gemma-3-27b-it",
            messages=[
                {
                    "role": "user",
                    "content": [
                    {
                        "type": "text",
                        "text": "Describe this image in markdown format"
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
                        }
                    }
                ]
            }
        ],
        )

        return completion.choices[0].message.content

class WebSearchPerplexityTool(Tool):
    name = "web_search_perplexity_tool"
    description = """This tool searches the web for information using Perplexity, try first to forward the entire question"""
    inputs = {
        "question": {"type": "string", "description": "Use the entire question at first and then iterate with the follow-up questions"}
    }
    output_type = "string"

    def forward(self, question: str) -> str:
        messages = [
            {
                "role": "system",
                "content": (
                    "Answer the question based on the information provided and following the instructions"
                ),
            },
            {   
                "role": "user",
                "content": question,
            },
        ]

        client = OpenAI(api_key=os.getenv("PERPLEXITY_API_KEY"), base_url="https://api.perplexity.ai")

        # chat completion without streaming
        response = client.chat.completions.create(
            model="sonar-pro",
            messages=messages,
        )
        
        return response.choices[0].message.content


def final_answer_formatting(answer, question):
    model = LiteLLMModel(
        model_id="gemini/gemini-2.0-flash",
        api_key=os.getenv("GOOGLE_API_KEY"),
    )

    prompt = f"""
    You are an AI assistant specialized in the GAIA benchmark. For the question provided, generate the answer in the exact format requested by the question. Do not include any other text, formatting or creative additions.

    For example:
    Question: "what is the capital of France ?"
    Answer: "The capital of France is Paris. It is the largest city in France and serves as the country's political, cultural, and economic center. Paris has been the capital since its liberation in 1944, although it has held this status intermittently throughout French history"
    Return: "Paris"

    Your turn:
    Question: {question}
    Answer: {answer}
    """

    messages = [
    {"role": "user", "content": [{"type": "text", "text": prompt}]}
    ]

    output = model(messages).content
    return output


web_agent = CodeAgent(
    model=LiteLLMModel(
        model_id="gemini/gemini-2.0-flash",
        api_key=os.getenv("GOOGLE_API_KEY"),
    ),
    tools=[
        WikipediaSearchTool(),
        GoogleSearchTool(provider="serper"), 
        VisitWebpageTool()
    ],
    add_base_tools=False,
    additional_authorized_imports=[
        "os", "requests", "inspect", "pandas", 
        "datetime", "re", "bs4", "markdownify"
    ],
    max_steps=10,
    name="web_agent",
    description="This agent is used to search the web for information"
)

audio_agent = CodeAgent(
    model=LiteLLMModel(
        model_id="gemini/gemini-2.0-flash",
        api_key=os.getenv("GOOGLE_API_KEY"),
    ),
    tools=[AudioToTextTool()],
    add_base_tools=False,
    max_steps=10,
    name="audio_agent",
    description="This agent is used to analyze and transcribe audio files"
)


manager_agent = CodeAgent(
    name="manager_agent",
    model=LiteLLMModel(
        model_id="gemini/gemini-2.0-flash",
        api_key=os.getenv("GOOGLE_API_KEY"),
    ),
    tools=[GetFileTool(), LoadXlsxFileTool(), LoadTextFileTool(), ImageAnalysisTool(), WebSearchPerplexityTool()],
    managed_agents=[audio_agent],
    additional_authorized_imports=[
        "pandas"
    ],
    planning_interval=5,
    verbosity_level=1,
    max_steps=10,
)