agent.py

import os
import pathlib

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

from prompt import calc_agent_prompt, web_search_agent_prompt

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

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

if not OPENAI_API_KEY and not ANTHROPIC_API_KEY:
    raise ValueError(
        "Please set either OPENAI_API_KEY or ANTHROPIC_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 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:
        return self.agent.run(question)


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}")