Class Tools

Toolsets and Stateful Tools

It's straightforward to create sets of related tools, and stateful tools by using the provided @toolset decorator on a class. Simply:

make a class containing all the functionality and methods you want the LLM to be able to access
decorate the methods the LLM can use with @tool(). Undecorated methods will be hidden from the LLM
pass either the class constructor or an instance of the class with the list of tools given to the LLM agent

from archytas.tool_utils import tool

class MyToolset:
    """
    general explanation of the class and how it's used
    """
    def __init__(self, ...):
        ... #initialize, set member values, etc.
    
    @tool
    def tool1(self, a:int) -> bool:
        """
        Simple description of what tool1 does

        More detailed description of the tool. This can be multiple lines.
        Explain what this tool does, how it's used, how it relates to other tools in the toolset. 
        Do not use python syntax in the description (since the LLM doesn't call functions that way).

        Args:
            a (int): Description of a

        Returns:
            bool: Description of the return value
        """
        ... #implementation

    @tool
    def tool2(self, s:str, v:dict) -> float:
        """
        <full docstring for this method>
        """
        ... #implementation

    @tool
    def tool3(self, ...) -> type:
        """
        <full docstring for this method>
        """
        ... #implementation

Methods should be type annotated, and contain matching docstrings, same as with function tools (noting that you should not mention the self parameter in the docstring).

SIR Toolset Example

Here's an example toolset for managing/running SIR simulations, adjusting simulation parameters, etc.

from archytas.tool_utils import tool


class ModelSimulation:
    """
    Simple example of a SIR model simulation
    """
    def __init__(self, dt=0.1):
        self._default_parameters = {'beta': 0.002, 'gamma': 0.1, 'S': 990, 'I': 10, 'R': 0}
        self.parameters = self._default_parameters.copy()
        self.dt = dt

    @tool
    def get_model_parameters(self) -> dict:
        """
        Get the model parameters

        Returns:
            dict: The model parameters in the form {param0: value0, param1: value1, ...}

        """
        return self.parameters

    @tool
    def set_model_parameters(self, update:dict):
        """
        Set some or all of the model parameters

        Args:
            update (dict): The parameters to update. Should be a dict of the form {param0: value0, param1: value1, ...}. Only the parameters specified will be updated.
        """
        self.parameters.update(update)

    @tool
    def run_model(self, steps:int=100) -> dict:
        """
        Run the model for a number of steps

        Args:
            steps (int): The number of steps to run the model for. Defaults to 100.

        Returns:
            dict: The model results in the form {param0: value0, param1: value1, ...}
        """
        S_new, I_new, R_new = self.parameters['S'], self.parameters['I'], self.parameters['R']
        beta, gamma = self.parameters['beta'], self.parameters['gamma']
        population = S_new + I_new + R_new

        for _ in range(steps):
            S_old, I_old, R_old = S_new, I_new, R_new
            dS = -beta * S_old * I_old
            dI = beta * S_old * I_old - gamma * I_old
            dR = gamma * I_old

            S_new = max(0, min(S_old + self.dt*dS, population))
            I_new = max(0, min(I_old + self.dt*dI, population))
            R_new = max(0, min(R_old + self.dt*dR, population))

            # Ensure the total population remains constant
            total_error = population - (S_new + I_new + R_new)
            R_new += total_error

        self.parameters['S'], self.parameters['I'], self.parameters['R'] = S_new, I_new, R_new
        return self.parameters

    @tool
    def reset_model(self):
        """
        Reset the model to the initial parameters
        """
        self.parameters = self._default_parameters.copy()

Notice that the class contains persistent data, as well as multiple methods that the LLM could call.

Passing the toolset to the agent

When instantiating an agent with this toolset, or any toolset, you have the option to pass the class constructor directly, or create your own instance which you pass in.

# directly pass the class constructor to tools
tools = [ModelSimulation]
agent = ReActAgent(tools=tools, verbose=True)

This is convenient if the default instantiation of the class is adequate for your use case. To use this approach, your class must be able to be instantiated with zero arguments (i.e. all arguments have a default value)

# create an instance and modify before passing in to agent
sim = ModelSimulation(dt=0.05)
sim._default_parameters = {'beta': 0.003, 'gamma': 0.2, 'S': 9990, 'I': 10, 'R': 0}
sim.reset_model()

tools = [sim]
agent = ReActAgent(tools=tools, verbose=True)

this is useful if you want to override the default initialization of the class.

@tool decorator

When a @tool decorator is applied to any method on a class, the following happens:

Create the prompt for the LLM for the whole toolset and all its contained tools.

For the model simulation tool above, the following prompt gets generated:

ModelSimulation (class):
    Simple example of a SIR model simulation

    methods:
        get_model_parameters:
            Get the model parameters

            _input_: None
            _output_: (dict) The model parameters in the form {param0: value0, param1: value1, ...}

        reset_model:
            Reset the model to the initial parameters

            _input_: None
            _output_: None
        
        run_model:
            Run the model for a number of steps

            _input_: (int, optional) The number of steps to run the model for. Defaults to 100.
            _output_: (dict) The model results in the form {param0: value0, param1: value1, ...}

        set_model_parameters:
            Set some or all of the model parameters

            _input_: a json object with the following fields:
            {
                "update": # (dict) The parameters to update. Should be a dict of the form {param0: value0, param1: value1, ...}. Only the parameters specified will be updated.
            }
            _output_: None

managing calling the methods, and passing results back to the LLM
- this process is very similar to the way it's handled for function tools. The main difference is that internally, the Archytas keeps an instance of the class, and routes the calls to the correct method based on the LLM's selected actions.
- The LLM refers to tools in a toolset via dot notation, so for the ModelSimulation toolset, the LLM can refer to the following tools:
  - ModelSimulation.get_model_parameters
  - ModelSimulation.run_model
  - ModelSimulation.set_model_parameters
  - ModelSimulation.reset_model
If the user asked the LLM to e.g. set the model parameter beta=0.004, the LLM might make the following call to the toolset:
```
{
    "thought": "The user wants me to set the beta parameter to 0.004",
    "tool": "ModelSimulation.set_model_parameters",
    "tool_input": {"beta": 0.004}
}
```