grady / agent.py
bstraehle's picture
Update agent.py
d8d4f52 verified
raw
history blame
829 Bytes
from smolagents import (
CodeAgent,
OpenAIServerModel,
DuckDuckGoSearchTool,
WikipediaSearchTool
)
def run_agent(model_id, level, question, file_name, ground_truth):
if level == 1:
max_steps = 5
elif level == 2:
max_steps = 10
elif level == 3:
max_steps = 20
print("###")
print(level)
print(max_steps)
print("###")
agent = CodeAgent(
add_base_tools = True,
additional_authorized_imports = ["pandas", "numpy"],
max_steps = max_steps,
model = OpenAIServerModel(model_id = model_id),
tools = [DuckDuckGoSearchTool(),
WikipediaSearchTool()],
verbosity_level = 2
)
answer = agent.run(question)
return answer, str(answer == ground_truth)