Creating Tools

In CodeCompanion, tools offer pre-defined ways for LLMs to execute actions on your machine, acting as an Agent in the process. This guide walks you through the implementation of tools, enabling you to create your own.

In the plugin, tools work by sharing a system prompt with an LLM. This instructs the LLM how to produce an XML markdown code block which can, in turn, be interpreted by the plugin to execute a command or function.

Architecture

In order to create tools, you do not need to understand the underlying architecture. However, for those who are curious about the implementation, please see the diagram below:

Building Your First Tool

Before we begin, it's important to familiarise yourself with the directory structure of the agents and tools implementation:

strategies/chat/agents
├── init.lua
├── executor/
│   ├── cmd.lua
│   ├── func.lua
│   ├── init.lua
│   ├── queue.lua
├── tools/
│   ├── cmd_runner.lua
│   ├── editor.lua
│   ├── files.lua
│   ├── rag.lua

When a tool is detected, the chat buffer sends any output to the agents/init.lua file (I will commonly refer to that as the "agent file" throughout this document). The agent file then parses the output into XML, identifying the tool and duly executing it.

There are two types of tools that CodeCompanion can leverage:

1. Command-based: These tools can execute a series of commands in the background using a plenary.job. They're non-blocking, meaning you can carry out other activities in Neovim whilst they run. Useful for heavy/time-consuming tasks.
2. Function-based: These tools, like @editor, execute Lua functions directly in Neovim within the main process, one after another.

For the purposes of this section of the guide, we'll be building a simple function-based calculator tool that an LLM can use to do basic maths.

Tool Structure

All tools must implement the following structure which the bulk of this guide will focus on explaining:

---@class CodeCompanion.Agent.Tool
---@field name string The name of the tool
---@field cmds table The commands to execute
---@field schema table The schema that the LLM must use in its response to execute a tool
---@field system_prompt fun(schema: table, xml2lua: table): string The system prompt to the LLM explaining the tool and the schema
---@field opts? table The options for the tool
---@field env? fun(schema: table): table|nil Any environment variables that can be used in the *_cmd fields. Receives the parsed schema from the LLM
---@field handlers table Functions which handle the execution of a tool
---@field handlers.setup? fun(agent: CodeCompanion.Agent): any Function used to setup the tool. Called before any commands
---@field handlers.on_exit? fun(agent: CodeCompanion.Agent): any Function to call at the end of a group of commands or functions
---@field output? table Functions which handle the output after every execution of a tool
---@field output.prompt fun(agent: CodeCompanion.Agent, self: CodeCompanion.Agent.Tool): string The message which is shared with the user when asking for their approval
---@field output.rejected? fun(agent: CodeCompanion.Agent, cmd: table): any Function to call if the user rejects running a command
---@field output.error? fun(agent: CodeCompanion.Agent, cmd: table, stderr: table, stdout?: table): any The function to call if an error occurs
---@field output.success? fun(agent: CodeCompanion.Agent, cmd: table, stdout: table): any Function to call if the tool is successful
---@field request table The request from the LLM to use the Tool

cmds

Command-Based Tools

The cmds table is a collection of commands which the agent will execute one after another, asynchronously, using plenary.job.

cmds = {
  { "make", "test" },
  { "echo", "hello" },
}

In this example, the plugin will execute make test followed by echo hello. After each command executes, the plugin will automatically send the output to a corresponding table on the agent file. If the command ran with success the output will be written to stdout, otherwise it will go to stderr. We'll be covering how you access that data in the output section below.

It's also possible to pass in environment variables (from the env function) by use of ${} brackets. The now removed @code_runner tool used them as below:

cmds = {
    { "docker", "pull", "${lang}" },
    {
      "docker",
      "run",
      "--rm",
      "-v",
      "${temp_dir}:${temp_dir}",
      "${lang}",
      "${lang}",
      "${temp_input}",
    },
  },
},
---@param xml table The values from the XML returned by the LLM
env = function(xml)
  local temp_input = vim.fn.tempname()
  local temp_dir = temp_input:match("(.*/)")
  local lang = xml.lang
  local code = xml.code

  return {
    code = code,
    lang = lang,
    temp_dir = temp_dir,
    temp_input = temp_input,
  }
end,

IMPORTANT

Using the handlers.setup() function, it's also possible to create commands dynamically like in the @cmd_runner tool.

Function-based Tools

Function-based tools use the cmds table to define functions that will be executed one after another. Each function has four parameters, itself, the actions request by the LLM, any input from a previous function call and a output_handler callback for async execution. The output_handler handles the result for an asynchronous tool. For a synchronous tool (like the calculator) you can ignore it. For the purpose of our calculator example:

cmds = {
  ---@param self CodeCompanion.Agent.Tool The Tools object
  ---@param actions table The action object
  ---@param input? any The output from the previous function call
  ---@param output_handler fun(output: { status: "success"|"error", data: any })? callback for async tool
  ---@return nil|{ status: "success"|"error", data: any }
  function(self, actions, input, output_handler)
-- Get the numbers and operation requested by the LLM
    local num1 = tonumber(actions.num1)
    local num2 = tonumber(actions.num2)
    local operation = actions.operation

    -- Validate input
    if not num1 then
      return { status = "error", data = "First number is missing or invalid" }
    end

    if not num2 then
      return { status = "error", data = "Second number is missing or invalid" }
    end

    if not operation then
      return { status = "error", data = "Operation is missing" }
    end

    -- Perform the calculation
    local result
    if operation == "add" then
      result = num1 + num2
    elseif operation == "subtract" then
      result = num1 - num2
    elseif operation == "multiply" then
      result = num1 * num2
    elseif operation == "divide" then
      if num2 == 0 then
        return { status = "error", data = "Cannot divide by zero" }
      end
      result = num1 / num2
    else
      return { status = "error", data = "Invalid operation: must be add, subtract, multiply, or divide" }
    end

    return { status = "success", data = result }
  end,
},

For a synchronous tool, you only need to return the result table as demonstrated. However, if you need to invoke some asynchronous actions in the tool, you can use the output_handler to submit any results to the executor, which will then invoke output functions to handle the results:

cmds = {
  function(agent, actions, input, output_handler)
    -- this is for demonstration only
    vim.lsp.client.request(lsp_method, lsp_param, function(err, result, _, _)
      agent.chat:add_message({role = "user", content = vim.json.encode(result)})
      output_handler({status = "success", data = result})
    end, buf_nr)
  end
}

Note that:

1. the output_handler will be called only once. Subsequent calls will be discarded;
2. A tool function should EITHER return the result table (synchronous), OR call the output_handler with the result table as the only argument (asynchronous), but not both. If a function tries to both return the result and call the output_handler, the result will be undefined because there's no guarantee which output will be handled first.

Similarly with command-based tools, the output is written to the stdout or stderr tables on the agent file. However, with function-based tools, the user must manually specify the outcome of the execution which in turn redirects the output to the correct table:

return { status = "error", data = "Invalid operation: must be add, subtract, multiply, or divide" }

Will cause execution of the tool to stop and populate stderr on the agent file.

return { status = "success", data = result }

Will populate the stdout table on the agent file and allow for execution to continue.

schema

The XML that the LLM has sent, is parsed and sent to the actions parameter of any function you've created in cmds, as a Lua table. If the LLM has done its job correctly, the Lua table should be the representation of what you've described in the schema.

In summary, the schema represents the structure of the response that the LLM must follow in order to call the tool. The schema is structured in Lua before being converted into XML with the awesome xml2lua parser.

For our basic calculator tool, which does an operation on two numbers, the schema could look something like:

schema = {
  _attr = { name = "calculator" },
  action = {
    num1 = "100",
    num2 = "50",
    operation = "multiply"
  },
},

system_prompt

In the plugin, LLMs are given knowledge about a tool and its schema via a system prompt. This method also informs the LLM on how to use the tool to achieve a desired outcome.

For our calculator tool, our system_prompt could look something like:

---@param schema table
---@param xml2lua table
---@return string
system_prompt = function(schema, xml2lua)
  return string.format([[## Calculator Tool (`calculator`)

### Purpose:
- To do a mathematical operation on two numbers.

### When to Use:
- Only invoke the calculator tool when the user specifically asks you

### Execution Format:
- Always return an XML markdown code block

### XML Schema:
Each tool invocation should adhere to this structure:

```xml
%s
```

where:
- `num1` is the first number to do any calculations with
- `num2` is the second number to do any calculations with
- `operation` is the mathematical operation to do on the two numbers. It MUST be one of `add`, `subtract`, `multiply` or `divide`

### Reminder:
- Minimize extra explanations and focus on returning correct XML blocks.
- Always use the structure above for consistency.]],
    xml2lua.toXml(schema, "tool")
  )
end,

You'll notice that the system_prompt function has two parameters:

• The schema table that we created earlier
• The xml2lua library that comes with the plugin

Whilst the latter is not mandated for use in the system prompt, it's the most optimum way of converting your schema to XML. You can see in the example above how we're using it. As a result, the system prompt that will be sent to the LLM will be:

## Calculator Tool (`calculator`) - Enhanced Guidelines

### Purpose:
- To do a mathematical operation on two numbers.

### When to Use:
- Only invoke the calculator tool when the user specifically asks you

### Execution Format:
- Always return an XML markdown code block

### XML Schema:
Each tool invocation should adhere to this structure:

```xml
<tool name="calculator">
  <action>
    <num1>100</num1>
    <num2>50</num2>
    <operation>multiply</operation>
  </action>
</tool>
```

where:
- `num1` is the first number to do any calculations with
- `num2` is the second number to do any calculations with
- `operation` is the mathematical operation to do on the two numbers. It MUST be one of `add`, `subtract`, `multiply` or `divide`

### Reminder:
- Minimize extra explanations and focus on returning correct XML blocks.
- Always use the structure above for consistency.

handlers

The handlers table contains two functions that are executed before and after a tool completes:

1. setup - Is called before anything in the cmds and output table. This is useful if you wish to set the cmds dynamically on the tool itself, like in the @cmd_runner tool.
2. on_exit - Is called after everything in the cmds and output table.

For the purposes of our calculator, let's just return some notifications so you can see the agent and tool flow:

handlers = {
  setup = function(agent)
    return vim.notify("setup function called", vim.log.levels.INFO)
  end,
  on_exit = function(agent)
    return vim.notify("on_exit function called", vim.log.levels.INFO)
  end,
},

output

The output table enables you to manage and format output from the execution of the cmds. It contains four functions:

1. success - Is called after every successful execution of a command/function. This can be a useful way of notifying the LLM of the success.
2. error - Is called when an error occurs whilst executing a command/function. It will only ever be called once as the whole execution of the cmds is halted. This can be a useful way of notifying the LLM of the failure.
3. prompt - Is called when user approval to execute the cmds is required. It forms the message prompt which the user is asked to confirm or reject.
4. rejected - Is called when a user rejects the approval to run the cmds. This method is used to inform the LLM of the rejection.

Let's consider how me might implement this for our calculator tool:

output = {
  ---@param agent CodeCompanion.Agent
  ---@param cmd table The command that was executed
  ---@param stdout table
  success = function(agent, cmd, stdout)
    local config = require("codecompanion.config")
    return agent.chat:add_buf_message({
      role = config.constants.USER_ROLE,
      content = string.format("The output from the calculator was %d", tonumber(stdout[1]))
    })
  end,
  error = function(agent)
    return vim.notify("An error occurred", vim.log.levels.ERROR)
  end,
},

Running the Calculator tool

If we put this all together in our config:

require("codecompanion").setup({
  strategies = {
    chat = {
      tools = {
        calculator = {
          description = "Perform calculations",
          callback = {
            name = "calculator",
            cmds = {
              function(self, actions, input)
                -- Get the numbers and operation requested by the LLM
                local num1 = tonumber(actions.num1)
                local num2 = tonumber(actions.num2)
                local operation = actions.operation

                -- Validate input
                if not num1 then
                  return { status = "error", data = "First number is missing or invalid" }
                end

                if not num2 then
                  return { status = "error", data = "Second number is missing or invalid" }
                end

                if not operation then
                  return { status = "error", data = "Operation is missing" }
                end

                -- Perform the calculation
                local result
                if operation == "add" then
                  result = num1 + num2
                elseif operation == "subtract" then
                  result = num1 - num2
                elseif operation == "multiply" then
                  result = num1 * num2
                elseif operation == "divide" then
                  if num2 == 0 then
                    return { status = "error", data = "Cannot divide by zero" }
                  end
                  result = num1 / num2
                else
                  return {
                    status = "error",
                    data = "Invalid operation: must be add, subtract, multiply, or divide",
                  }
                end

                return { status = "success", data = result }
              end,
            },
            schema = {
              _attr = { name = "calculator" },
              action = {
                num1 = "100",
                num2 = "50",
                operation = "multiply",
              },
            },
            system_prompt = function(schema, xml2lua)
              return string.format(
                [[## Calculator Tool (`calculator`) - Enhanced Guidelines

### Purpose:
- To do a mathematical operation on two numbers.

### When to Use:
- Only invoke the calculator tool when the user specifically asks you

### Execution Format:
- Always return an XML markdown code block

### XML Schema:
Each tool invocation should adhere to this structure:

```xml
%s
```

where:
- `num1` is the first number to do any calculations on
- `num2` is the second number to do any calculations on
- `operation` is the mathematical operation to do on the two numbers. It MUST be one of `add`, `subtract`, `multiply` or `divide`

### Reminder:
- Minimize extra explanations and focus on returning correct XML blocks.
- Always use the structure above for consistency.]],
                xml2lua.toXml(schema, "tool")
              )
            end,
            handlers = {
              setup = function(agent)
                return vim.notify("setup function called", vim.log.levels.INFO)
              end,
              on_exit = function(agent)
                return vim.notify("on_exit function called", vim.log.levels.INFO)
              end,
            },
            output = {
              ---@param agent CodeCompanion.Agent
              ---@param cmd table The command that was executed
              ---@param stdout table
              success = function(agent, cmd, stdout)
                local config = require("codecompanion.config")
                return agent.chat:add_buf_message({
                  role = config.constants.USER_ROLE,
                  content = string.format("The output from the calculator was %d", tonumber(stdout[1])),
                })
              end,
              error = function(agent)
                return vim.notify("An error occurred", vim.log.levels.ERROR)
              end,
            },
          },
        },
      },
    }
  }
})

and with the prompt:

@calculator what is 100*50?

You should see: The output from the calculator was 5000, in the chat buffer.

Adding in User Approvals

A big concern for users when they create and deploy their own tools is "what if an LLM does something I'm not aware of or I don't approve?". To that end, CodeCompanion tries to make it easy for a user to be the "human in the loop" and approve tool use before execution.

To enable this for any tool, simply add the requires_approval = true in a tool's opts table:

require("codecompanion").setup({
  strategies = {
    chat = {
      tools = {
        calculator = {
          description = "Perform calculations",
          callback = "as above",
          opts = {
            requires_approval = true,
          },
        }
      }
    }
  }
})

To account for the user being prompted for an approval, we can add a output.prompt to the tool:

output = {
  -- success and error functions remain the same ...

  ---The message which is shared with the user when asking for their approval
  ---@param agent CodeCompanion.Agent
  ---@param self CodeCompanion.Agent.Tool
  ---@return string
  prompt = function(agent, self)
    return string.upper(self.request.action.operation) .. " " .. self.request.action.num1  .. " and " .. self.request.action.num2 .. "?"
  end,
},

This will notify the user with the message: MULTIPLY 100 and 50?. The user can choose to proceed, reject or cancel. The latter will cancel any tools from running.

You can also customize the output if a user rejects the approval:

output = {
  -- success, error and prompt functions remain the same ...

    ---Rejection message back to the LLM
    ---@param agent CodeCompanion.Agent
    ---@return nil
    rejected = function(agent)
      local config = require("codecompanion.config")
      return agent.chat:add_buf_message({
        role = config.constants.USER_ROLE,
        content = "The user rejected your request to run the calculator tool"
      })
  end,
},

request

At runtime, CodeCompanion will take the LLMs XML request and add that to a request table on the tool itself to make it easy to access.