File size: 10,936 Bytes
d8d14f1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
# ToolAgent Documentation

The `ToolAgent` class is a specialized agent that facilitates the execution of specific tasks using a model and tokenizer. It is part of the `swarms` module and inherits from the `Agent` class. This agent is designed to generate functions based on a given JSON schema and task, making it highly adaptable for various use cases, including natural language processing and data generation.

The `ToolAgent` class plays a crucial role in leveraging pre-trained models and tokenizers to automate tasks that require the interpretation and generation of structured data. By providing a flexible interface and robust error handling, it ensures smooth integration and efficient task execution.

### Parameters

| Parameter          | Type                              | Description                                                                     |
|--------------------|-----------------------------------|---------------------------------------------------------------------------------|
| `name`             | `str`                             | The name of the tool agent. Default is "Function Calling Agent".                |
| `description`      | `str`                             | A description of the tool agent. Default is "Generates a function based on the input json schema and the task". |
| `model`            | `Any`                             | The model used by the tool agent.                                               |
| `tokenizer`        | `Any`                             | The tokenizer used by the tool agent.                                           |
| `json_schema`      | `Any`                             | The JSON schema used by the tool agent.                                         |
| `max_number_tokens`| `int`                             | The maximum number of tokens for generation. Default is 500.                    |
| `parsing_function` | `Optional[Callable]`              | An optional parsing function to process the output of the tool agent.           |
| `llm`              | `Any`                             | An optional large language model to be used by the tool agent.                  |
| `*args`            | Variable length argument list     | Additional positional arguments.                                                |
| `**kwargs`         | Arbitrary keyword arguments       | Additional keyword arguments.                                                   |

### Attributes

| Attribute          | Type  | Description                                  |
|--------------------|-------|----------------------------------------------|
| `name`             | `str` | The name of the tool agent.                  |
| `description`      | `str` | A description of the tool agent.             |
| `model`            | `Any` | The model used by the tool agent.            |
| `tokenizer`        | `Any` | The tokenizer used by the tool agent.        |
| `json_schema`      | `Any` | The JSON schema used by the tool agent.      |

### Methods

#### `run`

```python
def run(self, task: str, *args, **kwargs) -> Any:
```

**Parameters:**

| Parameter  | Type                      | Description                                                      |
|------------|---------------------------|------------------------------------------------------------------|
| `task`     | `str`                     | The task to be performed by the tool agent.                      |
| `*args`    | Variable length argument list | Additional positional arguments.                                 |
| `**kwargs` | Arbitrary keyword arguments | Additional keyword arguments.                                    |

**Returns:** 

- The output of the tool agent.

**Raises:** 

- `Exception`: If an error occurs during the execution of the tool agent.

## Functionality and Usage

The `ToolAgent` class provides a structured way to perform tasks using a model and tokenizer. It initializes with essential parameters and attributes, and the `run` method facilitates the execution of the specified task.

### Initialization

The initialization of a `ToolAgent` involves specifying its name, description, model, tokenizer, JSON schema, maximum number of tokens, optional parsing function, and optional large language model.

```python
agent = ToolAgent(
    name="My Tool Agent",
    description="A tool agent for specific tasks",
    model=model,
    tokenizer=tokenizer,
    json_schema=json_schema,
    max_number_tokens=1000,
    parsing_function=my_parsing_function,
    llm=my_llm
)
```

### Running a Task

To execute a task using the `ToolAgent`, the `run` method is called with the task description and any additional arguments or keyword arguments.

```python
result = agent.run("Generate a person's information based on the given schema.")
print(result)
```

### Detailed Examples

#### Example 1: Basic Usage

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from swarms import ToolAgent

model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b")
tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")

json_schema = {
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "age": {"type": "number"},
        "is_student": {"type": "boolean"},
        "courses": {
            "type": "array",
            "items": {"type": "string"}
        }
    }
}

task = "Generate a person's information based on the following schema:"
agent = ToolAgent(model=model, tokenizer=tokenizer, json_schema=json_schema)
generated_data = agent.run(task)

print(generated_data)
```

#### Example 2: Using a Parsing Function

```python
def parse_output(output):
    # Custom parsing logic
    return output

agent = ToolAgent(
    name="Parsed Tool Agent",
    description="A tool agent with a parsing function",
    model=model,
    tokenizer=tokenizer,
    json_schema=json_schema,
    parsing_function=parse_output
)

task = "Generate a person's information with custom parsing:"
parsed_data = agent.run(task)

print(parsed_data)
```

#### Example 3: Specifying Maximum Number of Tokens

```python
agent = ToolAgent(
    name="Token Limited Tool Agent",
    description="A tool agent with a token limit",
    model=model,
    tokenizer=tokenizer,
    json_schema=json_schema,
    max_number_tokens=200
)

task = "Generate a concise person's information:"
limited_data = agent.run(task)

print(limited_data)
```


## Full Usage
```python

from pydantic import BaseModel, Field
from transformers import AutoModelForCausalLM, AutoTokenizer

from swarms import ToolAgent
from swarms.tools.json_utils import base_model_to_json

# Model name
model_name = "CohereForAI/c4ai-command-r-v01-4bit"

# Load the pre-trained model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
)

# Load the pre-trained model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)


# Initialize the schema for the person's information
class APIExampleRequestSchema(BaseModel):
    endpoint: str = Field(
        ..., description="The API endpoint for the example request"
    )
    method: str = Field(
        ..., description="The HTTP method for the example request"
    )
    headers: dict = Field(
        ..., description="The headers for the example request"
    )
    body: dict = Field(..., description="The body of the example request")
    response: dict = Field(
        ...,
        description="The expected response of the example request",
    )


# Convert the schema to a JSON string
api_example_schema = base_model_to_json(APIExampleRequestSchema)
# Convert the schema to a JSON string

# Define the task to generate a person's information
task = "Generate an example API request using this code:\n"

# Create an instance of the ToolAgent class
agent = ToolAgent(
    name="Command R Tool Agent",
    description=(
        "An agent that generates an API request using the Command R"
        " model."
    ),
    model=model,
    tokenizer=tokenizer,
    json_schema=api_example_schema,
)

# Run the agent to generate the person's information
generated_data = agent.run(task)

# Print the generated data
print(f"Generated data: {generated_data}")



```


## Jamba ++ ToolAgent
```python
from pydantic import BaseModel, Field
from transformers import AutoModelForCausalLM, AutoTokenizer

from swarms import ToolAgent
from swarms.tools.json_utils import base_model_to_json

# Model name
model_name = "ai21labs/Jamba-v0.1"

# Load the pre-trained model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
)

# Load the pre-trained model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)


# Initialize the schema for the person's information
class APIExampleRequestSchema(BaseModel):
    endpoint: str = Field(
        ..., description="The API endpoint for the example request"
    )
    method: str = Field(
        ..., description="The HTTP method for the example request"
    )
    headers: dict = Field(
        ..., description="The headers for the example request"
    )
    body: dict = Field(..., description="The body of the example request")
    response: dict = Field(
        ...,
        description="The expected response of the example request",
    )


# Convert the schema to a JSON string
api_example_schema = base_model_to_json(APIExampleRequestSchema)
# Convert the schema to a JSON string

# Define the task to generate a person's information
task = "Generate an example API request using this code:\n"

# Create an instance of the ToolAgent class
agent = ToolAgent(
    name="Command R Tool Agent",
    description=(
        "An agent that generates an API request using the Command R"
        " model."
    ),
    model=model,
    tokenizer=tokenizer,
    json_schema=api_example_schema,
)

# Run the agent to generate the person's information
generated_data = agent(task)

# Print the generated data
print(f"Generated data: {generated_data}")
```

## Additional Information and Tips

- Ensure that either the `model` or `llm` parameter is provided during initialization. If neither is provided, the `ToolAgent` will raise an exception.
- The `parsing_function` parameter is optional but can be very useful for post-processing the output of the tool agent.
- Adjust the `max_number_tokens` parameter to control the length of the generated output, depending on the requirements of the task.

## References and Resources

- [Transformers Documentation](https://huggingface.co/transformers/)
- [Loguru Logger](https://loguru.readthedocs.io/en/stable/)

This documentation provides a comprehensive guide to the `ToolAgent` class, including its initialization, usage, and practical examples. By following the detailed instructions and examples, developers can effectively utilize the `ToolAgent` for various tasks involving model and tokenizer-based operations.