grady / agent.py
bstraehle's picture
Update agent.py
0c53631 verified
raw
history blame
1.53 kB
import os
from openai import OpenAI
from smolagents import (
CodeAgent,
OpenAIServerModel,
DuckDuckGoSearchTool,
WikipediaSearchTool
)
MODEL_ID = "o4-mini"
def run_agent(level, question, file_name):
if level == 1:
max_steps = 10
elif level == 2:
max_steps = 25
elif level == 3:
max_steps = 50
agent = CodeAgent(
add_base_tools = True,
additional_authorized_imports = ["pandas", "numpy"],
max_steps = max_steps,
model = OpenAIServerModel(model_id = MODEL_ID),
#planning_interval=3,
tools = [#DuckDuckGoSearchTool(),
WikipediaSearchTool()],
verbosity_level = 1
)
agent.logger.console.width = 66
answer = agent.run(question)
return answer, get_final_answer(question, answer)
def get_final_answer(question, answer):
client = OpenAI()
prompt_template = """
You are a GAIA benchmark assistant, who is given a question and an answer. Read the question carefully. You must give the final answer requested. Do not include explanations, steps, reasoning, or additional text.
**Question:** """ + str(question) + """
**Answer:** """ + str(answer) + """
**Final answer:**:
"""
completion = client.chat.completions.create(
messages = [{"role": "user", "content": [{"type": "text", "text": prompt_template}]}],
model = MODEL_ID
)
final_answer = completion.choices[0].message.content
return final_answer