Update agent.py
Browse files
agent.py
CHANGED
@@ -5,15 +5,22 @@ from smolagents import (
|
|
5 |
WikipediaSearchTool
|
6 |
)
|
7 |
|
8 |
-
def run_agent(model_id, question):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
9 |
agent = CodeAgent(
|
|
|
|
|
|
|
10 |
model=OpenAIServerModel(model_id=model_id),
|
11 |
tools=[DuckDuckGoSearchTool(),
|
12 |
WikipediaSearchTool()],
|
13 |
-
add_base_tools=True,
|
14 |
-
max_steps=5,
|
15 |
verbosity_level=2,
|
16 |
-
additional_authorized_imports=["pandas", "numpy"]
|
17 |
)
|
18 |
|
19 |
return agent.run(question)
|
|
|
5 |
WikipediaSearchTool
|
6 |
)
|
7 |
|
8 |
+
def run_agent(model_id, level, question, file_name):
|
9 |
+
if level == 1:
|
10 |
+
max_steps = 5
|
11 |
+
elif level == 2:
|
12 |
+
max_steps = 10
|
13 |
+
elif level == 3:
|
14 |
+
max_steps = 20
|
15 |
+
|
16 |
agent = CodeAgent(
|
17 |
+
add_base_tools=True,
|
18 |
+
additional_authorized_imports=["pandas", "numpy"]
|
19 |
+
max_steps=max_steps,
|
20 |
model=OpenAIServerModel(model_id=model_id),
|
21 |
tools=[DuckDuckGoSearchTool(),
|
22 |
WikipediaSearchTool()],
|
|
|
|
|
23 |
verbosity_level=2,
|
|
|
24 |
)
|
25 |
|
26 |
return agent.run(question)
|