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
971 Bytes
import importlib
import yaml
from smolagents import CodeAgent
from tools import analyze_image, read_pdf
def create_media_agent(model):
"""
Create a specialized agent for handling media (images, PDFs).
Args:
model: The model to use for the agent
Returns:
Configured CodeAgent for media handling
"""
# Load default prompts
prompt_templates = yaml.safe_load(
importlib.resources.files("smolagents.prompts")
.joinpath("code_agent.yaml")
.read_text()
)
media_agent = CodeAgent(
tools=[analyze_image, read_pdf],
model=model,
name="media_agent",
description="Specialized agent for handling media files like images and PDFs. Use this agent to analyze images and extract text from PDF documents.",
add_base_tools=True,
additional_authorized_imports=["PIL", "io", "requests"],
prompt_templates=prompt_templates,
)
return media_agent