We’ll use Cohere’s Python SDK for the code examples. Follow along in this notebook.

Introduction

In this chapter, we’ll dissect the key components of a tool use system and what a tool use workflow looks like. And we’ll do that with a concrete code example.

To build an intuition of how tool use works, let’s start with a simple example. We’ll look at a use case of a RAG assistant that can query the sales database of an e-commerce company.

In summary, there are two parts to building a tool use application:

Setting up the tools
Running the tool use workflow

Setting Up the Tools

First, let’s install the Cohere Python SDK and set up the Cohere client.

# pip install cohere

import json
import cohere
co = cohere.Client(api_key="YOUR_COHERE_API_KEY") # Get your API key: https://dashboard.cohere.com/api-keys

The pre-requisite, or Step 0, before we can run a tool use workflow, is to set up the tools. We can break this further into two steps:

Creating the tool
Defining the tool schema

The pre-requisite of running a tool use workflow is to set up the tools

Step 1: Create the Tool

Let’s create a function to query a sales database called daily_sales_report and represent it as a tool. For simplicity, it contains a mock database containing just three data entries and the logic to return the data given a user query.

In this example, we are defining a Python function as the tool. But more broadly, the tool can be any function or service that can receive and send data. It could be an email service, an SQL database, a vector database, a weather data service, a sports data service, a web search engine, or even another LLM, just to give a few examples.

def daily_sales_report(day: str) -> dict:
    """
    Function to retrieve the sales report for the given day
    """
    # Mock database containing daily sales reports
    sales_database = {
    '2023-09-28': {'total_sales_amount': 5000,'total_units_sold': 100},
    '2023-09-29': {'total_sales_amount': 10000,'total_units_sold': 250},
    '2023-09-30': {'total_sales_amount': 8000,'total_units_sold': 200}
    }
    
    report = sales_database.get(day, {})
    
    if report:
        return {
            'date': day,
            'summary': f"Total Sales Amount: {report['total_sales_amount']}, Total Units Sold: {report['total_units_sold']}"
        }
    else:
        return {'date': day, 'summary': 'No sales data available for this day.'}
    

functions_map = {
    "daily_sales_report": daily_sales_report
}

Step 2: Define Tool Schema

Next, we define the tool schema for the sales database tool. This schema is what will be passed to the Cohere API when running a tool use workflow. The schema must contain the following fields:

name: the name of the tool.
description: a description of what the tool is and what it is used for.
parameter_definitions: a list of parameters that the tool accepts. For each parameter, we need to define the following fields:
- description (details about the parameter)
- type (the parameter’s data type, e.g., str, int, and so on)
- required (either True or False) fields

This schema informs the LLM about what the tool does, and the LLM decides whether to use a particular tool based on it. Therefore, the more descriptive and specific the schema, the more likely the LLM will make the right tool call decisions.

Let’s name our tool daily_sales_report and populate the schema. The function has only one parameter, day, which requires the caller to pass a date in the YYYY-MM-DD string format.

tools = [
    {
        "name": "daily_sales_report",
        "description": "Connects to a database to retrieve overall sales volumes and sales information for a given day.",
        "parameter_definitions": {
            "day": {
                "description": "Retrieves sales data for this day, formatted as YYYY-MM-DD.",
                "type": "str",
                "required": True
            }
        }
    }
]

Step 3: Create Custom Preamble (Optional)

An optional step is to add a custom preamble, which is the LLM’s system message.

The recommended approach is to use two H2 Markdown headers: "Task & Context" and "Style Guide" in the exact order (read the complete guide).

It’s a completely optional step, though it’s likely needed if we want to create a robust and reliable application. Also note that the preamble is not related to the tool setup that we covered earlier, rather it’s part of the instruction to the LLM.

Let’s create a custom preamble as follows:

preamble = """## Task & Context
You help people answer their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You should focus on serving the user's needs as best you can, which will be wide-ranging.

## Style Guide
Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling.
"""

Running the tool workflow

The tool is now ready to use. We can think of a tool use system as consisting of four components:

The user
The application
The LLM
The tools

At its most basic, these four components interact in a workflow through four steps:

Step 1: Get user message. The LLM gets the user message (via the application).
Step 2: Generate tool calls. The LLM decides which tools to call (if any) and generates the tool calls.
Step 3: Get tool results. The application executes the tools, and the results are sent to the LLM.
Step 4: Generate response and citations. The LLM generates the response and citations back to the user.

An overview of the four-step tool use workflow

Let’s walk through these four steps using a code example. In this first example, we’ll use the simplest possible scenario where:

There is only one tool (the sales database)
Tool calling happens only once (and only one tool is called)
There is only one turn in the conversation (no conversation memory preserved)

Step 1: Get User Message

Step 1: getting the user message — Step 1: Getting the user message

In the first step, the application receives the user's message. As an example, we’ll use “Can you provide a sales summary for 29th September 2023?”

message = "Can you provide a sales summary for 29th September 2023?"

Step 2: Generate Tool Calls

During this step:

The application calls the Chat endpoint to generate the list of tool calls. We add the following parameters:
- message - the user message.
- model - the model of choice. Here we use Command R+.
- preamble - the custom preamble.
- tools - the list of tools we defined earlier (only one in this example).
- force_single_step - we define it as True. This forces the endpoint to run tool calls only once and not over a sequence of runs (we’ll see how the opposite – multi-step – works in Chapter 4).
The endpoint sends back a list of tool calls to be made (if it determines that tools are required).

Here’s the code for running this step.

model = "command-r-plus"

# Initial response to the user message
response = co.chat(
    message=message,
    model=model,
    preamble=preamble,
    tools=tools,
    force_single_step=True
)

tool_calls = response.tool_calls
    
print("Tool calls:\n")
for i, t in enumerate(tool_calls):
    print(f"#{i+1}\nTool: {t.name}\nParameters: {t.parameters}\n")

And here’s the list of tool calls (only one in our case). The model determines that to answer the user question, it needs to call the daily_sales_report tool with the parameter 'day': '2023-09-29' – as we’d expect.

Tool calls:

#1
Tool: daily_sales_report
Parameters: {'day': '2023-09-29'}

Step 3: Get Tool Results

During this step, the application performs the function calling. It calls the necessary tools based on the tool call payloads given by the endpoint. In our case, it calls the daily_sales_report function.

Here’s the code for running this step.

tool_results = []
for tc in tool_calls:
    tool_output = functions_map[tc.name](**tc.parameters)
    tool_call = {"name": tc.name, "parameters": tc.parameters}
    tool_results.append({"call": tool_call, "outputs": [tool_output]})
    
print("Tool results:\n")
for i, t in enumerate(tool_results):
    print(f"#{i+1}\nTool call: {t['call']}\nOutputs: {t['outputs']}\n")

And here are the tool results (again, only one in our case). The tool successfully returns the correct data on the sales amount and units sold.

Tool results:

#1
Tool call: {'name': 'daily_sales_report', 'parameters': {'day': '2023-09-29'}}
Outputs: [{'date': '2023-09-29', 'summary': 'Total Sales Amount: 10000, Total Units Sold: 250'}]

Step 4: Generate Response and Citations

Step 4: Generating the response and citations

We have now come to the fourth and final step: response generation. By this time, the tool call has already been executed, and the result has been returned to the LLM.

During this step, the application calls the Chat endpoint to generate the response to the user. Compared to the call in Step 2, here are the parameter changes required:

message - we put a blank value "" as this is a call to generate the chatbot’s response rather than to get a new user message (which doesn’t change)
tool_results - here we introduce this parameter as a way to pass the tool results generated in the previous step
chat_history - we use this parameter to pass the chat_history value from Step 2’s response. It contains the context required to generate the response – the user message and the tool call(s) made – as shown below.

# response.chat_history

[Message_User(message='Can you provide a sales summary for 29th September 2023?', tool_calls=None, role='USER'),
 Message_Chatbot(tool_calls=[ToolCall(name='daily_sales_report', parameters={'day': '2023-09-29'})], role='CHATBOT')]

Here’s the code for running this step:

# Generate response
response = co.chat(
    message="",
    model=model,
    preamble=preamble,
    tools=tools,
    tool_results=tool_results,
    chat_history=response.chat_history
)
    
print("Final response:")
print(response.text)
print("="*50)

if response.citations:
    print("\nCitations:")
    for citation in response.citations:
        print(citation)
    print("\nCited Documents:")
    for document in response.documents:
        print(document)

And here’s the final response from the LLM, which accurately answers the user’s question (total sales amount: 10,000, and total number of units sold: 250).

It also generates fine-grained citations, which are included with Command R+ out of the box. Here, we see the model generating two citations, one for each specific span in its response, where it uses the tool result to answer the question.

For example, the first citation shows that the total sales amount was 10,000 portion of the response comes from the daily_sales_report tool. It also provides the values of the start and end character numbers that can be used for front-end development purposes to render citations easily and elegantly.

Final response:
On 29 September 2023, we sold 250 units, totalling 10,000 in sales.
==================================================

Citations:
start=30 end=39 text='250 units' document_ids=['daily_sales_report:0:2:0']
start=51 end=66 text='10,000 in sales' document_ids=['daily_sales_report:0:2:0']

Cited Documents:
{'date': '2023-09-29', 'id': 'daily_sales_report:0:2:0', 'summary': 'Total Sales Amount: 10000, Total Units Sold: 250', 'tool_name': 'daily_sales_report'}

The screenshot below, taken from Cohere’s Chat UI, shows an example of displaying fine-grained citations on a user interface. In this example, a web search tool is used to find the capital of Ontario. The response includes highlighting Toronto as the span that cites the web search results.

An example of displaying fine-grained citations on a user interface, taken from Cohere’s Chat UI

Streaming the Response

We can stream the response from the Chat endpoint for each generated token instead of having to wait for the full response. To enable streaming, we need to change the endpoint call from co.chat to co.chat_stream.

The code from Step 4 would then change to the following. The difference now is we are capturing the response events as the endpoint streams them.

# Generate response
response = co.chat_stream(
    message="",
    model=model,
    preamble=preamble,
    tools=tools,
    tool_results=tool_results,
    chat_history=response.chat_history
)
    
print("Final response:\n")
chatbot_response = ""

for event in response:
    if event.event_type == "text-generation":
        print(event.text, end="")
        chatbot_response += event.text
    if event.event_type == "stream-end":
        if event.response.citations:
            print("\n\nCitations:")
            for citation in event.response.citations:
                print(citation)
        if event.response.documents:
            print("\nCited Documents")
            for document in event.response.documents:
                print(document)

Tool Use Modes

A good tool use API must be able to handle different scenarios, from simple to complex LLM tool calls. This translates to being able to call tools in different modes: single-step, multi-step, and parallel.

The good news is that Cohere’s API supports all these modes. We’ll walk through them in detail in the coming two chapters, but for now, let’s briefly look at what these different modes are.

Cohere’s API supports different modes of tool use

Single-Step

This is the most straightforward mode, so let's start here. An LLM calls a single tool, the tool returns the result, and the result gets appended to the prompt.

Referring to the four-step workflow above, this means that Steps 2 (generating the tool calls) and 3 (getting the tool results) happen exactly once.

This is the scenario that we saw in the sales database example earlier.

To illustrate this further, let's assume we have two tools:

A web search tool: It takes a query and returns relevant search results from the web
A Python interpreter: It takes a block of Python code, runs it, and returns the output

Suppose a user asks, "What was Apple's revenue in 2023?"

To answer this, a single web search of "Apple 2023 revenue" will suffice.

Single-Step, Parallel

The LLM can also decide that more than one tool call is required, which means calling tools in parallel. This can be calling the same tool multiple times or different tools for any number of calls.

We’ll see an example of this in Chapter 3.

Suppose the user now asks, "What were Apple's and Google's revenue in 2023?"

To answer this, two web searches of "Apple 2023 revenue" and "Google 2023 revenue" are needed.

Multi-Step

This is needed when more than one tool is required, and the output of one tool is needed as the input to the other tool.

In other words, tool-calling needs to happen in a sequence.

Referring to the four-step workflow above, this means that Steps 2 (generating the tool calls) and 3 (getting the tool results) can happen multiple times before moving on to Step 4 (generating the response).

We’ll see examples of this in Chapters 4 and 5.

Suppose the user now asks, "Plot Apple's 2023 revenue on a bar chart".

This requires first a web search for Apple's revenue, and then using the search result to generate the Python code for plotting.

Multi-Step, Parallel

Tool calling can run in parallel in a multi-step scenario as it can in the single-step case.

We’ll see examples of this in Chapters 4 and 5.

Let's say the user now asks, "Plot Apple's and Google's 2023 revenue on a bar chart".

This requires two parallel web search calls — one each for Apple and Google — and then running the Python plotting code in the second step.

Conclusion

In this chapter, we looked at the components of a tool use system, the tool use workflow, state management in tool use, and different modes of tool use.

We also looked at a code example using a single-step scenario. But it was a basic example that doesn’t address some questions like:

How to set up a system with more than one tool?
Does the assistant always need to call tools?
How does the assistant work in a chat setting consisting of multiple turns?

In Chapter 3, we’ll address these questions and improve on the code example in this chapter.

Contents