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
decorate the whole class with @toolset()
pass either the class constructor or an instance of the class with the list of tools given to the LLM agent

@toolset()
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

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 toolset, tool


@toolset()
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

    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.

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.01)
sim._default_parameters = {'beta': 0.003, 'gamma': 0.2, 'S': 9990, 'I': 10, 'R': 0}
sim.reset()

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