agent.py

import os
import pathlib

import dotenv
from smolagents import (
    CodeAgent,
    OpenAIServerModel,
    VisitWebpageTool,
    WebSearchTool,
    WikipediaSearchTool,
)

from prompt import (
    calc_agent_prompt,
    gen_GAIA_answer_formatter_prompt,
    web_search_agent_prompt,
)

if pathlib.Path(".env").exists():
    dotenv.load_dotenv(".env")

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", None)

if not OPENAI_API_KEY:
    raise ValueError(
        "Please set either OPENAI_API_KEY in your .env file or Huggingface Space settings."
    )

model = OpenAIServerModel(
    model_id="gpt-4.1",
    api_key=OPENAI_API_KEY,
)

web_search_agent = CodeAgent(
    model=model,
    name="web_search_agent",
    description="An agent that can search the web, visit webpages, and summarize information.",
    prompt_templates=web_search_agent_prompt,
    additional_authorized_imports=["requests", "beautifulsoup4"],
    tools=[
        WebSearchTool(),
        VisitWebpageTool(),
        WikipediaSearchTool(),
    ],
)

calc_agent = CodeAgent(
    model=model,
    name="calc_agent",
    description="An agent that can perform calculations, solve math problems, and analyze data.",
    prompt_templates=calc_agent_prompt,
    additional_authorized_imports=["pandas", "numpy", "math", "statistics", "scipy"],
    tools=[]
)

class ExtractFailedException(Exception):
    """Custom exception for failed extraction of formatted answer."""
    pass

class GAIAAnswerFormatter:

    model = OpenAIServerModel(
        model_id="gpt-4.1-mini",
        api_key=OPENAI_API_KEY,
    )

    def extract_formatted_answer(self, llm_response: str) -> str:
        import re 

        match = re.search(r'<formated_answer>(.*?)</formated_answer>', llm_response, re.DOTALL)
        if match:
            return match.group(1).strip()
        raise ExtractFailedException(
            "Failed to extract formatted answer from the LLM response."
        )

    def __call__(self, question: str, answer: str) -> str:
        message = [{
            "role": "user",
            "content": gen_GAIA_answer_formatter_prompt(question, answer),
        }]
        response = self.model.generate(messages=message)
        try:
            return self.extract_formatted_answer(response.content)
        except ExtractFailedException:
            return answer


answer_formatter = GAIAAnswerFormatter()

class Agent:
    def __init__(self):
        self.agent = CodeAgent(
            model=model,
            managed_agents=[web_search_agent, calc_agent],
            tools=[],
            add_base_tools=True,
        )


    def __call__(self, question: str) -> str:
        answer = self.agent.run(question)
        return answer_formatter(question, answer)


if __name__ == "__main__":
    GAIA_EXAMPLE_QUESTION = """
In NASA’s Astronomy Picture of the Day on 2006 January 21, two astronauts are visible,
with one appearing much smaller than the other. As of August 2023, out of the astronauts in the
NASA Astronaut Group that the smaller astronaut was a member of, which one spent the least time
in space, and how many minutes did he spend in space, rounded to the nearest minute? Exclude any
astronauts who did not spend any time in space. Give the last name of the astronaut, separated from
the number of minutes by a semicolon. Use commas as thousands separators in the number of minutes.
"""
    agent = Agent()
    final_answer = agent(GAIA_EXAMPLE_QUESTION)

    print(f"\n\n---\n\nFinal answer: {final_answer}")