Dependency Injection

Sometimes when writing a tool, it is useful to have access to the agent, or aspects of the archytas environment. Tools are written as functions/classes to be called by the LLM, so in order to get access, we provide a dependency injection framework.

The following dependencies are available to be hooked into:

agent: (AgentRef) the underlying agent with connections to the OpenAI API
tool function: (ToolFnRef) the instance of the function being called by the agent
tool name: (ToolNameRef) the name of the function being called by the agent
loop controller: (LoopControllerRef) an object for breaking out of the react loop early

To use a dependency, simply include an argument in your @tool function signature with the type of the dependency you want. e.g.

@tool
def myTool(args..., agent:AgentRef):
    ...

@tool
def myTool(args..., raw_tool:ToolFnRef):
    ...

@tool
def myTool(args..., tool_name:ToolNameRef):
    ...

@tool
def myTool(args..., loop_controller:LoopControllerRef):
    ...

Or more than one at a time if you need multiple

@tool
def myTool(args..., agent:AgentRef, raw_tool:ToolFnRef, tool_name:ToolNameRef, loop_controller:LoopControllerRef):
    ...

Note that the argument name need not exactly match, but it is good practice to use the standard names for injected dependencies.

Agent

A reference to the underlying LLM agent. This can be useful in a number of cases. It allows your tool to access anything that the agent has such as chat history. You can also use the agent to make sub-queries to the LLM.

The pirate_subquery demo tool demonstrates receiving the injected agent reference, and using it to make a sub-query to the same agent

@tool()
def pirate_subquery(query:str, agent:AgentRef) -> str:
    """
    Runs a subquery using a oneshot agent in which answers will be worded like a pirate.

    Args:
        query (str): The query to run against the agent.

    Returns:
        str: Result of the subquery in pirate vernacular.

    """
    prompt = """
    You are an pirate. Answer all questions truthfully using pirate vernacular.
    """
    return agent.oneshot(prompt=prompt, query=query)

agent:AgentRef is a special type that is ignored when generating the prompt. Instead at runtime, when the function is called, it gets passed a reference to the agent, which is used in the body of the function.

Tool Function

This is the instance of the function (or instance method) that is being called by the LLM. This is relatively unexplored in terms of usefulness.

Tool Name

This is just the name of the tool being called, as provided by the LLM. Typically this will exactly match the current function name, unless the name was overridden

Loop Controller

This is an object mainly for managing early exits from the react loop. E.g. you have a tool that you want the agent to call, and then after it should always exit the react loop without considering calling other tools. This dependency allows for both graceful and ungraceful exits from the react loop.

An example use case is the dataset toolset

@toolset()
class DatasetToolset:

    dataset_id: Optional[int]
    df: Optional[pd.DataFrame]

    def __init__(self, *args, **kwargs): ...

    def set_dataset(self, dataset_id, agent=None): ...

    def load_dataframe(self, filename=None): ...

    def reset(self): ...

    def send_dataset(self): ...

    def context(self):
        return f"""You are an analyst whose goal is to help with scientific data analysis and manipulation in Python.

You are working on a dataset named: {self.dataset.get('name')}

The description of the dataset is:
{self.dataset.get('description')}

The dataset has the following structure:
--- START ---
{self.dataset_info()}
--- END ---

Please answer any user queries to the best of your ability, but do not guess if you are not sure of an answer.
If you are asked to manipulate or visualize the dataset, use the generate_python_code tool.
"""

    @tool()
    def dataset_info(self) -> str:
        """
        Inspect the dataset and return information and metadata about it.

        This should be used to answer questions about the dataset, including information about the columns,
        and default parameter values and initial states.


        Returns:
            str: a textual representation of the dataset
        """
        # Need to actually track things. Maybe a good idea to split this in to finer tools so certain things can be queried?
        output = f"""
Dataframe head:
{self.df.head(15)}
Columns:
{self.df.columns}
dtypes:
{self.df.dtypes}
Statistics:
{self.df.describe()}
"""
        return output

    @tool()
    def generate_python_code(
        self, query: str, agent: AgentRef, loop: LoopControllerRef
    ) -> str:
        """
        Generated Python code to be run in an interactive Jupyter notebook for the purpose of exploring, modifying and visualizing a Pandas Dataframe.

        Input is a full grammatically correct question about or request for an action to be performed on the loaded dataframe.

        Assume that the dataframe is already loaded and has the variable name `df`.
        Information about the dataframe can be loaded with the `dataset_info` tool.

        Args:
            query (str): A fully grammatically correct queistion about the current dataset.

        Returns:
            str: A LLM prompt that should be passed evaluated.
        """
        # set up the agent
        prompt = ...
        llm_response = agent.oneshot(prompt=prompt, query=query)
        loop.set_state(loop.STOP_SUCCESS)
        preamble, code, coda = re.split("```\w*", llm_response)
        result = json.dumps(
            {
                "action": "code_cell",
                "language": "python",
                "content": code.strip(),
            }
        )
        return result