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

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.

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.

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

Detailed Examples

Example 1: Basic Usage

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

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

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

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

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
• Loguru Logger

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.