File size: 3,501 Bytes
7194b60 989f731 8baf998 7194b60 682a11a 7194b60 682a11a 7194b60 8baf998 7194b60 8baf998 7194b60 989f731 7194b60 8baf998 7194b60 989f731 8baf998 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | 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}") |