User-Defined Agent Tools

Define custom tools for the RAG Agent

There are many cases where it is helpful to define custom tools for the RAG Agent. R2R allows for users to define custom tools, passing these definitions into the Agent at server start.

Defining New Tools

There is a directory in the R2R repository, /docker/user_tools, which is mounted to the R2R docker container. It is here that we will place our custom tool files.

There, we will find a README.md file, which includes a template for our new tool:

1from core.base.agent.tools.base import Tool
2
3
4class ToolNameTool(Tool):
5 """
6 A user defined tool.
7 """
8
9 def __init__(self):
10 super().__init__(
11 name="tool_name",
12 description="A natural language tool description that is shown to the agent.",
13 parameters={
14 "type": "object",
15 "properties": {
16 "input_parameter": {
17 "type": "string",
18 "description": "Define any input parameters by their name and type",
19 },
20 },
21 "required": ["input_parameter"],
22 },
23 results_function=self.execute,
24 llm_format_function=None,
25 )
26
27 async def execute(self, input_parameter: str, *args, **kwargs):
28 """
29 Implementation of the tool.
30 """
31
32 # Any custom tool logic can go here
33
34 output_response = some_method(input_parameter)
35
36 result = AggregateSearchResult(
37 generic_tool_result=[web_response],
38 )
39
40 # Add to results collector if context is provided
41 if context and hasattr(context, "search_results_collector"):
42 context.search_results_collector.add_aggregate_result(result)
43
44 return result

This template has two basic methods:

  1. __init__ is where we define the tool. The description that we make here is shown to the agent.
  2. execute is where we define any custom tool logic and interact with the inputs.

Writing our new tool

Below, we have an example of a toy tool, which takes an integer and string input, returning a silly message to the agent. Should your tool require additional dependencies, be sure to include them in the user_requirements.txt file located in the /docker directory.

1from r2r import Tool, AggregateSearchResult
2
3
4class SecretMethodTool(Tool):
5 """
6 A user defined tool.
7 """
8
9 def __init__(self):
10 super().__init__(
11 name="secret_method",
12 description="Performs a secret method.",
13 parameters={
14 "type": "object",
15 "properties": {
16 "number": {
17 "type": "string",
18 "description": "An integer input for the secret method.",
19 },
20 "string": {
21 "type": "string",
22 "description": "A string input for the secret method.",
23 },
24 },
25 "required": ["number", "string"],
26 },
27 results_function=self.execute,
28 llm_format_function=None,
29 )
30
31 async def execute(self, number: int, string: str, *args, **kwargs):
32 """
33 Implementation of the tool.
34 """
35
36 output_response = f"Your order for {number} dancing flamingos has been received. They will arrive by unicycle courier within 3-5 business dreams. Please prepare {string} for them."
37
38 result = AggregateSearchResult(
39 generic_tool_result=output_response,
40 )
41
42 context = self.context
43 # Add to results collector if context is provided
44 if context and hasattr(context, "search_results_collector"):
45 context.search_results_collector.add_aggregate_result(result)
46
47 return result

Finally, we can modify our configuration file’s agent section to include our new tool:

1[agent]
2rag_tools = ["secret_method"]

Finally, we can run the following and see that our agent called our new method, passed the required parameters, and understood its output:

1client.retrieval.agent(
2 message={"role": "user", "content": "Can you run the secret method tool? Feel free to use any parameters you want. I just want to see the output."},
3)