Spaces:
Sleeping
Sleeping
# 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. |