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
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