Performing Tasks Sequentially with Cohere's RAG

Compare two user queries to a RAG chatbot, “What was Apple’s revenue in 2023?” and “What was the revenue of the most valuable company in the US in 2023?”.

While the first query is straightforward to handle, the second query requires breaking down into two steps:

1. Identify the most valuable company in the US in 2023
2. Get the revenue of the company in 2023

These steps need to happen in a sequence rather than all at once. This is because the information retrieved from the first step is required to inform the second step.

This is an example of sequential reasoning. In this tutorial, we’ll learn how agentic RAG with Cohere handles sequential reasoning, and in particular:

• Multi-step tool calling
• Multi-step, parallel tool calling
• Self-correction

We’ll learn these by building an agent that answers questions about using Cohere.

Setup

To get started, first we need to install the cohere library and create a Cohere client.

We also need to import the tool definitions that we’ll use in this tutorial.

! pip install cohere langchain langchain-community pydantic -qq

import json
import os
import cohere
from tool_def import (
    search_developer_docs,
    search_developer_docs_tool,
    search_internet,
    search_internet_tool,
    search_code_examples,
    search_code_examples_tool,
)
co = cohere.ClientV2(
    "COHERE_API_KEY"
)  # Get your free API key: https://dashboard.cohere.com/api-keys
os.environ["TAVILY_API_KEY"] = (
    "TAVILY_API_KEY"  # We'll need the Tavily API key to perform internet search. Get your API key: https://app.tavily.com/home
)

Setting up the tools

We set up the same set of tools as in Part 1. If you want further details on how to set up the tools, check out Part 1.

functions_map = {
    "search_developer_docs": search_developer_docs,
    "search_internet": search_internet,
    "search_code_examples": search_code_examples,
}

Running an agentic RAG workflow

We create a run_agent function to run the agentic RAG workflow, the same as in Part 1. If you want further details on how to set up the tools, check out Part 1.

tools = [
    search_developer_docs_tool,
    search_internet_tool,
    search_code_examples_tool,
]

system_message = """## Task and Context
You are an assistant who helps developers use Cohere. You are equipped with a number of tools that can provide different types of information. If you can't find the information you need from one tool, you should try other tools if there is a possibility that they could provide the information you need."""

model = "command-a-plus-05-2026"
def run_agent(query, messages=None):
    if messages is None:
        messages = []
    if "system" not in {m.get("role") for m in messages}:
        messages.append({"role": "system", "content": system_message})
    # Step 1: get user message
    print(f"QUESTION:\n{query}")
    print("=" * 50)
    messages.append({"role": "user", "content": query})
    # Step 2: Generate tool calls (if any)
    response = co.chat(
        model=model, messages=messages, tools=tools, temperature=0.3
    )
    while response.message.tool_calls:
        print("TOOL PLAN:")
        print(response.message.tool_plan, "\n")
        print("TOOL CALLS:")
        for tc in response.message.tool_calls:
            print(
                f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}"
            )
        print("=" * 50)
        messages.append(response.message)
        # Step 3: Get tool results
        for tc in response.message.tool_calls:
            tool_result = functions_map[tc.function.name](
                **json.loads(tc.function.arguments)
            )
            tool_content = []
            for data in tool_result:
                tool_content.append(
                    {
                        "type": "document",
                        "document": {"data": json.dumps(data)},
                    }
                )
                # Optional: add an "id" field in the "document" object, otherwise IDs are auto-generated
            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": tool_content,
                }
            )
        # Step 4: Generate response and citations
        response = co.chat(
            model=model,
            messages=messages,
            tools=tools,
            temperature=0.3,
        )
    messages.append(
        {
            "role": "assistant",
            "content": response.message.content[0].text,
        }
    )
    # Print final response
    print("RESPONSE:")
    print(response.message.content[0].text)
    print("=" * 50)
    # Print citations (if any)
    verbose_source = (
        False  # Change to True to display the contents of a source
    )
    if response.message.citations:
        print("CITATIONS:\n")
        for citation in response.message.citations:
            print(
                f"Start: {citation.start}| End:{citation.end}| Text:'{citation.text}' "
            )
            print("Sources:")
            for idx, source in enumerate(citation.sources):
                print(f"{idx+1}. {source.id}")
                if verbose_source:
                    print(f"{source.tool_output}")
            print("\n")
    return messages

Multi-step tool calling

Let’s ask the agent a few questions, starting with this one about a specific feature. The user is asking about two things: a feature to reorder search results and code examples for that feature.

In this case, the agent first needs to identify what that feature is before it can answer the second part of the question.

This is reflected in the agent’s tool plan, which describes the steps it will take to answer the question.

So, it first calls the search_developer_docs tool to find the feature.

It then discovers that the feature is Rerank. Using this information, it calls the search_code_examples tool to find code examples for that feature.

Finally, it uses the retrieved information to answer both parts of the user’s question.

messages = run_agent(
    "What's the Cohere feature to reorder search results? Do you have any code examples on that?"
)

QUESTION:
What's the Cohere feature to reorder search results? Do you have any code examples on that?
==================================================
TOOL PLAN:
I will search for the Cohere feature to reorder search results. Then I will search for code examples on that. 
TOOL CALLS:
Tool name: search_developer_docs | Parameters: {"query":"reorder search results"}
==================================================
TOOL PLAN:
I found that the Rerank endpoint is the feature that reorders search results. I will now search for code examples on that. 
TOOL CALLS:
Tool name: search_code_examples | Parameters: {"query":"rerank endpoint"}
==================================================
RESPONSE:
The Rerank endpoint is the feature that reorders search results. Unfortunately, I could not find any code examples on that.
==================================================
CITATIONS:
Start: 4| End:19| Text:'Rerank endpoint' 
Sources:
1. search_developer_docs_53tfk9zgwgzt:0

Multi-step, parallel tool calling

In Part 2, we saw how the Cohere API supports parallel tool calling, and in this tutorial, we looked at sequential tool calling. That also means that both scenarios can happen at the same time.

Here’s an example. Suppose we ask the agent to find the CEOs of the companies with the top 3 highest market capitalization.

In the first step, it searches the Internet for information about the 3 companies with the highest market capitalization.

And in the second step, it performs parallel searches for the CEOs of the 3 identified companies.

messages = run_agent(
    "Who are the CEOs of the companies with the top 3 highest market capitalization."
)

QUESTION:
Who are the CEOs of the companies with the top 3 highest market capitalization.
==================================================
TOOL PLAN:
I will search for the top 3 companies with the highest market capitalization. Then, I will search for the CEOs of those companies. 
TOOL CALLS:
Tool name: search_internet | Parameters: {"query":"top 3 companies with highest market capitalization"}
==================================================
TOOL PLAN:
The top 3 companies with the highest market capitalization are Apple, Microsoft, and Nvidia. I will now search for the CEOs of these companies. 
TOOL CALLS:
Tool name: search_internet | Parameters: {"query":"Apple CEO"}
Tool name: search_internet | Parameters: {"query":"Microsoft CEO"}
Tool name: search_internet | Parameters: {"query":"Nvidia CEO"}
==================================================
RESPONSE:
The CEOs of the top 3 companies with the highest market capitalization are:
1. Tim Cook of Apple
2. Satya Nadella of Microsoft
3. Jensen Huang of Nvidia
==================================================
CITATIONS:
Start: 79| End:87| Text:'Tim Cook' 
Sources:
1. search_internet_0f8wyxfc3hmn:0
2. search_internet_0f8wyxfc3hmn:1
3. search_internet_0f8wyxfc3hmn:2
Start: 91| End:96| Text:'Apple' 
Sources:
1. search_internet_kb9qgs1ps69e:0
Start: 100| End:113| Text:'Satya Nadella' 
Sources:
1. search_internet_wy4mn7286a88:0
2. search_internet_wy4mn7286a88:1
3. search_internet_wy4mn7286a88:2
Start: 117| End:126| Text:'Microsoft' 
Sources:
1. search_internet_kb9qgs1ps69e:0
Start: 130| End:142| Text:'Jensen Huang' 
Sources:
1. search_internet_q9ahz81npfqz:0
2. search_internet_q9ahz81npfqz:1
3. search_internet_q9ahz81npfqz:2
4. search_internet_q9ahz81npfqz:3
Start: 146| End:152| Text:'Nvidia' 
Sources:
1. search_internet_kb9qgs1ps69e:0

Self-correction

The concept of sequential reasoning is useful in a broader sense, particularly where the agent needs to adapt and change its plan midway in a task.

In other words, it allows the agent to self-correct.

To illustrate this, let’s look at an example. Here, the user is asking about the authors of the sentence BERT paper.

The agent attempted to find required information via the search_developer_docs tool.

However, we know that the tool doesn’t contain this information because we have only added a small sample of documents.

As a result, the agent, having received the documents back without any relevant information, decides to search the internet instead. This is also helped by the fact that we have added specific instructions in the search_internet tool to search the internet for information not found in the developer documentation.

It finally has the information it needs, and uses it to answer the user’s question.

This highlights another important aspect of agentic RAG, which allows a RAG system to be flexible. This is achieved by powering the retrieval component with an LLM.

On the other hand, a standard RAG system would typically hand-engineer this, and hence, is more rigid.

messages = run_agent(
    "Who are the authors of the sentence BERT paper?"
)

QUESTION:
Who are the authors of the sentence BERT paper?
==================================================
TOOL PLAN:
I will search for the authors of the sentence BERT paper. 
TOOL CALLS:
Tool name: search_developer_docs | Parameters: {"query":"authors of the sentence BERT paper"}
==================================================
TOOL PLAN:
I was unable to find any information about the authors of the sentence BERT paper. I will now search for 'sentence BERT paper authors'. 
TOOL CALLS:
Tool name: search_internet | Parameters: {"query":"sentence BERT paper authors"}
==================================================
RESPONSE:
The authors of the Sentence-BERT paper are Nils Reimers and Iryna Gurevych.
==================================================
CITATIONS:
Start: 43| End:55| Text:'Nils Reimers' 
Sources:
1. search_internet_z8t19852my9q:0
2. search_internet_z8t19852my9q:1
3. search_internet_z8t19852my9q:2
4. search_internet_z8t19852my9q:3
5. search_internet_z8t19852my9q:4
Start: 60| End:75| Text:'Iryna Gurevych.' 
Sources:
1. search_internet_z8t19852my9q:0
2. search_internet_z8t19852my9q:1
3. search_internet_z8t19852my9q:2
4. search_internet_z8t19852my9q:3
5. search_internet_z8t19852my9q:4

Summary

In this tutorial, we learned about:

• How multi-step tool calling works
• How multi-step, parallel tool calling works
• How multi-step tool calling enables an agent to self-correct, and hence, be more flexible

However, up until now, we have only worked with purely unstructured data, the type of data we typically encounter in a standard RAG system.

In the coming chapters, we’ll add another complexity to the agentic RAG system – working with semi-structured and structured data. This adds another dimension to the agent’s flexibility, which is dealing with a more diverse set of data sources.

In Part 4, we’ll learn how to build an agent that can perform faceted queries over semi-structured data.