The ToolAgent class is a specialized agent within the Zytron module, designed to facilitate the execution of specific tasks by utilizing a model and tokenizer. Inheriting from the base Agent class, this versatile agent is crafted to generate functions dynamically based on a provided JSON schema and task, making it highly adaptable for various use cases, including natural language processing, data generation, and workflow automation.
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 zytron 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 zytron import ToolAgent
from zytron.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 zytron import ToolAgent
from zytron.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.
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.
