File size: 984 Bytes
837e221
e4c7240
837e221
 
e4c7240
837e221
 
e4c7240
837e221
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e4c7240
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
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