Skip to content

Dependency Injection

David-Andrew Samson edited this page Aug 26, 2024 · 7 revisions

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

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: ...

    @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 a sub agent and have it generate the code 
        prompt = ... # tell the llm to write code to solve the query the user asked about the data
        llm_response = agent.oneshot(prompt=prompt, query=query)

        # Make agent exit react loop after calling this tool
        loop.set_state(loop.STOP_SUCCESS)
        
        # extract the generated code
        preamble, code, coda = re.split("```\w*", llm_response)

        # package up the code and return it for use
        result = ...
        return result

In this use case, the main agent uses a sub agent to generate, but not run, code which then gets used later in a jupyter notebook interface. There are other tools in this interface that the main agent may be used via the ReAct pattern, but when the LLM decides it needs some code, the system needs to break out of the react loop so the generated code can be handled.

To break out of the react loop, the generate_python_code function has an argument loop:LoopControllerRef which is then used in the body to indicate that the ReAct loop should be broken out of:

loop.set_state(loop.STOP_SUCCESS)

In addition to breaking out of the loop normally, the loop reference can be used to catastrophically break out of the ReAct loop, similar to how too many exceptions will end the loop early.

loop.set_state(loop.STOP_FATAL)

This will also stop the loop, and raise an exception at the call to agent.react() which may or may not be caught and then shown to the user, e.g. as in the example repl chat

A more in depth explanation of flow control and exception handling is in Flow Control and Errors

(TODO)

  • effect handler