agents/media_agent/agent.py

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