agent_config.py

from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, VisitWebpageTool, PythonInterpreterTool, LiteLLMModel, tool
from tools.wiki_tool import WikiTool
from tools.excel_analysis_tool import ExcelAnalysisTool
import os
import requests
import re
import tempfile
from pathlib import Path
from typing import Optional

@tool
def check_answer(answer: str) -> str:
    """
    Reviews the answer to check that it meets the requirements specified by the user and modifies it if necessary.
    Args:
        answer (str): The answer fo the Agent.
    Returns:
        str: The final answer.
    """
    return answer

@tool
def reverse_sentence_tool(reverse_sentence: str) -> str:
    """
    This tool receives a sentence where both the word order and the characters in each word are reversed.
    It returns the sentence with words and order corrected.
    Args:
        reverse_sentence: A sentence with reversed words and reversed word order.
    Returns:
        A sentence in natural reading order.
    """
    inverted_words = reverse_sentence.split(" ")[::-1]
    correct_words = [word[::-1] for word in inverted_words]

    return " ".join(correct_words)

@tool
def download_file_if_any(base_api_url: str, task_id: str) -> Optional[str]:
    """
    Downloads a file from a remote API using a task ID.

    Args:
        base_api_url (str): Base URL of the API to fetch files from (e.g., "https://example.com/api").
        task_id (str): Task identifier used to locate the file.

    Returns:
        Optional[str]: Path to the downloaded file, or None if the file is not found (404).
    """
    url = f"{base_api_url}/files/{task_id}"
    try:
        resp = requests.get(url, timeout=30)
        if resp.status_code == 404:
            return None  # file not found
        resp.raise_for_status()  # raise other errors
    except requests.exceptions.HTTPError as e:
        raise e  # propagate for logging/debugging

    # Extract filename from Content-Disposition if present
    cdisp = resp.headers.get("content-disposition", "")
    filename = task_id  # fallback
    if "filename=" in cdisp:
        match = re.search(r'filename="([^"]+)"', cdisp)
        if match:
            filename = match.group(1)

    tmp_dir = Path(tempfile.gettempdir()) / "gaia_files"
    tmp_dir.mkdir(parents=True, exist_ok=True)
    file_path = tmp_dir / filename

    with open(file_path, "wb") as f:
        f.write(resp.content)

    return str(file_path)

def create_agent():
    """
    Creates and configures the CodeAgent with the necessary model and tools.

    Returns:
        CodeAgent: The configured agent instance.
    """
    # Load API key from environment variable (more secure)
    token = os.getenv("HF_API_KEY")
    if not token:
        raise RuntimeError("Missing Hugging Face API key. Set HF_API_KEY environment variable.")

    # Initialize the LLM model
    # model = HfApiModel(token=token, model_id="gpt-3.5-turbo", provider="openai")
    
    model = LiteLLMModel(
            token=token,
            model_id="gemini/gemini-2.0-flash",
            api_key=os.getenv("GEMINI_API_KEY"),
        )

    # Initialize tools
    wiki_tool = WikiTool()
    web_search_tool = DuckDuckGoSearchTool()
    excel_analysis_tool = ExcelAnalysisTool()
    visitWebpageTool = VisitWebpageTool()
    python_interpreter_tool = PythonInterpreterTool()

    # Create and return the agent
    agent = CodeAgent(
        model=model,
        tools=[wiki_tool, web_search_tool, excel_analysis_tool, visitWebpageTool, python_interpreter_tool, check_answer, reverse_sentence_tool, download_file_if_any],
        add_base_tools=True,
        max_steps=8,
        verbosity_level=2,
        additional_authorized_imports=['numpy','subprocess', 're', 'pandas', 'requests', 'json', 'os', 'pathlib', 'tempfile'],
    )
    return agent

# Example usage
if __name__ == "__main__":
    agent = create_agent()
    print("Agent created with tools:", [tool.name for tool in agent.tools])