File size: 3,111 Bytes
5c0be56
 
 
55ef143
 
5c0be56
 
 
 
 
55ef143
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5c0be56
 
 
 
 
 
 
 
55ef143
 
 
 
 
 
 
5c0be56
55ef143
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5c0be56
 
 
 
 
 
 
55ef143
 
 
 
 
 
 
5c0be56
55ef143
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5c0be56
 
 
 
 
 
 
55ef143
 
 
 
 
 
 
5c0be56
55ef143
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import importlib

import yaml
from smolagents import CodeAgent

from prompts import (
    DATA_AGENT_SYSTEM_PROMPT,
    MEDIA_AGENT_SYSTEM_PROMPT,
    WEB_AGENT_SYSTEM_PROMPT,
)
from tools import (
    analyze_image,
    browse_webpage,
    extract_dates,
    find_in_page,
    parse_csv,
    perform_calculation,
    read_pdf,
    web_search,
)


def create_web_agent(model):
    """
    Create a specialized agent for web browsing tasks.

    Args:
        model: The model to use for the agent

    Returns:
        Configured CodeAgent for web browsing
    """

    prompt_templates = yaml.safe_load(
        importlib.resources.files("smolagents.prompts")
        .joinpath("code_agent.yaml")
        .read_text()
    )
    # prompt_templates["system_prompt"] = WEB_AGENT_SYSTEM_PROMPT

    web_agent = CodeAgent(
        tools=[web_search, browse_webpage, find_in_page, extract_dates],
        model=model,
        name="web_agent",
        description="Specialized agent for web browsing and searching. Use this agent to find information online, browse websites, and extract information from web pages.",
        add_base_tools=True,
        additional_authorized_imports=["requests", "bs4", "re", "json"],
        prompt_templates=prompt_templates,
    )

    return web_agent


def create_data_analysis_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
    """
    prompt_templates = yaml.safe_load(
        importlib.resources.files("smolagents.prompts")
        .joinpath("code_agent.yaml")
        .read_text()
    )
    # prompt_templates["system_prompt"] = DATA_AGENT_SYSTEM_PROMPT

    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


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
    """
    prompt_templates = yaml.safe_load(
        importlib.resources.files("smolagents.prompts")
        .joinpath("code_agent.yaml")
        .read_text()
    )
    # prompt_templates["system_prompt"] = MEDIA_AGENT_SYSTEM_PROMPT

    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