mjschock's picture
Enhance agent functionality in main_v2.py by adding WikipediaSearchTool and updating DuckDuckGoSearchTool and VisitWebpageTool parameters. Modify agent initialization to accommodate new tools and increase max results and output length. Update requirements.txt to include Wikipedia-API dependency. Refactor imports for better organization across agent modules.
e4c7240 unverified
raw
history blame contribute delete
984 Bytes
import importlib
import yaml
from smolagents import CodeAgent
from tools import parse_csv, perform_calculation
def create_data_agent(model):
"""
Create a specialized agent for data analysis tasks.
Args:
model: The model to use for the agent
Returns:
Configured CodeAgent for data analysis
"""
# Load default prompts
prompt_templates = yaml.safe_load(
importlib.resources.files("smolagents.prompts")
.joinpath("code_agent.yaml")
.read_text()
)
data_agent = CodeAgent(
tools=[parse_csv, perform_calculation],
model=model,
name="data_agent",
description="Specialized agent for data analysis. Use this agent to analyze data, perform calculations, and extract insights from structured data.",
add_base_tools=True,
additional_authorized_imports=["pandas", "numpy", "math", "csv", "io"],
prompt_templates=prompt_templates,
)
return data_agent