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

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)