# Routing Queries to Data Sources

> Build an agentic RAG system that routes queries to the most relevant tools based on the query's nature.

<a target="_blank" href="https://colab.research.google.com/github/cohere-ai/cohere-developer-experience/blob/main/notebooks/guides/agentic-rag/agentic_rag_pt1_routing.ipynb">
  Open in Colab
</a>

Imagine a RAG system that can search over diverse sources, such as a website, a database, and a set of documents.

In a standard RAG setting, the application would aggregate retrieved documents from all the different sources it is connected to. This may contribute to noise from less relevant documents.

Additionally, it doesn’t take into consideration that, given a data source's nature, it might be less or more relevant to a query than the other data sources.

An agentic RAG system can solve this problem by routing queries to the most relevant tools based on the query's nature. This is done by leveraging the tool use capabilities of the Chat endpoint.

In this tutorial, we'll cover:

* Setting up the tools
* Running an agentic RAG workflow
* Routing queries to tools

We'll build an agent that can answer questions about using Cohere, equipped with a number of different tools.

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

<Note>
   Important: the source code for tool definitions can be 

  [found here](https://github.com/cohere-ai/cohere-developer-experience/blob/main/notebooks/guides/agentic-rag/tool_def.py)

  . Make sure to have the 

  `tool_def.py`

   file in the same directory as this notebook for the imports to work correctly. 
</Note>

```python PYTHON
! pip install cohere langchain langchain-community pydantic -qq
```

```python PYTHON
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

In an agentic RAG system, each data source is represented as a tool. A tool is broadly any function or service that can receive and send objects to the LLM. But in the case of RAG, this becomes a more specific case of a tool that takes a query as input and returns a set of documents.

Here, we are defining a Python function for each tool, but more broadly, the tool can be any function or service that can receive and send objects.

* `search_developer_docs`: Searches Cohere developer documentation. Here we are creating a small list of sample documents for simplicity and will return the same list for every query. In practice, you will want to implement a search function such as those that use semantic search.
* `search_internet`: Performs an internet search using Tavily search, which we take from LangChain's ready implementation.
* `search_code_examples`: Searches for Cohere code examples and tutorials. Here we are also creating a small list of sample documents for simplicity.

These functions are mapped to a dictionary called `functions_map` for easy access.

Here, we are defining a Python function for each tool.

Further reading:

* [Documentation on parameter types in tool use](https://docs.cohere.com/v2/docs/parameter-types-in-tool-use)

```python PYTHON
functions_map = {
    "search_developer_docs": search_developer_docs,
    "search_internet": search_internet,
    "search_code_examples": search_code_examples,
}
```

The second and final setup step is to define the tool schemas in a format that can be passed to the Chat endpoint. A tool schema must contain the following fields: `name`, `description`, and `parameters` in the format shown below.

This schema informs the LLM about what the tool does, which enables an LLM to decide whether to use a particular tool. Therefore, the more descriptive and specific the schema, the more likely the LLM will make the right tool call decisions.

## Running an agentic RAG workflow

We can now run an agentic RAG workflow using a tool use approach. We can think of the 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: Tool planning and calling** – The LLM makes a decision on the tools to call (if any) and generates the tool calls
* **Step 3: Tool execution** - The application executes the tools and sends the results to the LLM
* **Step 4: Response and citation generation** – The LLM generates the response and citations to back to the user

We wrap all these steps in a function called `run_agent`.

```python PYTHON
tools = [
    search_developer_docs_tool,
    search_internet_tool,
    search_code_examples_tool,
]
```

```python PYTHON
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."""
```

```python PYTHON
model = "command-a-03-2025"


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
```

## Routing queries to tools

Let's ask the agent a few questions, starting with this one about the Embed endpoint.

Because the question asks about a specific feature, the agent decides to use the `search_developer_docs` tool (instead of retrieving from all the data sources it's connected to).

It first generates a tool plan that describes how it will handle the query. Then, it generates tool calls to the `search_developer_docs` tool with the associated `query` parameter.

The tool does indeed contain the information asked by the user, which the agent then uses to generate its response.

```python PYTHON
messages = run_agent("How many languages does Embed support?")
```

```mdx
QUESTION:
How many languages does Embed support?
==================================================
TOOL PLAN:
I will search the Cohere developer documentation for 'how many languages does Embed support'. 

TOOL CALLS:
Tool name: search_developer_docs | Parameters: {"query":"how many languages does Embed support"}
==================================================
RESPONSE:
The Embed endpoint supports over 100 languages.
==================================================
CITATIONS:

Start: 28| End:47| Text:'over 100 languages.' 
Sources:
1. search_developer_docs_gwt5g55gjc3w:2
```

Let's now ask the agent a question about setting up the Notion API so we can connect it to LLMs. This information is not likely to be found in the developer documentation or code examples because it is not Cohere-specific, so we can expect the agent to use the internet search tool.

And this is exactly what the agent does. This time, it decides to use the `search_internet` tool, triggers the search through Tavily search, and uses the results to generate its response.

```python PYTHON
messages = run_agent("How to set up the Notion API.")
```

```mdx
QUESTION:
How to set up the Notion API.
==================================================
TOOL PLAN:
I will search for 'Notion API setup' to find out how to set up the Notion API. 

TOOL CALLS:
Tool name: search_internet | Parameters: {"query":"Notion API setup"}
==================================================
RESPONSE:
To set up the Notion API, you need to create a new integration in Notion's integrations dashboard. You can do this by navigating to https://www.notion.com/my-integrations and clicking '+ New integration'.

Once you've done this, you'll need to get your API secret by visiting the Configuration tab. You should keep your API secret just that – a secret! You can refresh your secret if you accidentally expose it.

Next, you'll need to give your integration page permissions. To do this, you'll need to pick or create a Notion page, then click on the ... More menu in the top-right corner of the page. Scroll down to + Add Connections, then search for your integration and select it. You'll then need to confirm the integration can access the page and all of its child pages.

If your API requests are failing, you should confirm you have given the integration permission to the page you are trying to update.

You can also create a Notion API integration and get your internal integration token. You'll then need to create a .env file and add environmental variables, get your Notion database ID and add your integration to your database.

For more information on what you can build with Notion's API, you can refer to this guide.
==================================================
CITATIONS:

Start: 38| End:62| Text:'create a new integration' 
Sources:
1. search_internet_cwabyfc5mn8c:0
2. search_internet_cwabyfc5mn8c:2


Start: 75| End:98| Text:'integrations dashboard.' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 132| End:170| Text:'https://www.notion.com/my-integrations' 
Sources:
1. search_internet_cwabyfc5mn8c:0


Start: 184| End:203| Text:''+ New integration'' 
Sources:
1. search_internet_cwabyfc5mn8c:0
2. search_internet_cwabyfc5mn8c:2


Start: 244| End:263| Text:'get your API secret' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 280| End:298| Text:'Configuration tab.' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 310| End:351| Text:'keep your API secret just that – a secret' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 361| End:411| Text:'refresh your secret if you accidentally expose it.' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 434| End:473| Text:'give your integration page permissions.' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 501| End:529| Text:'pick or create a Notion page' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 536| End:599| Text:'click on the ... More menu in the top-right corner of the page.' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 600| End:632| Text:'Scroll down to + Add Connections' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 639| End:681| Text:'search for your integration and select it.' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 702| End:773| Text:'confirm the integration can access the page and all of its child pages.' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 783| End:807| Text:'API requests are failing' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 820| End:907| Text:'confirm you have given the integration permission to the page you are trying to update.' 
Sources:
1. search_internet_cwabyfc5mn8c:2


Start: 922| End:953| Text:'create a Notion API integration' 
Sources:
1. search_internet_cwabyfc5mn8c:1


Start: 958| End:994| Text:'get your internal integration token.' 
Sources:
1. search_internet_cwabyfc5mn8c:1


Start: 1015| End:1065| Text:'create a .env file and add environmental variables' 
Sources:
1. search_internet_cwabyfc5mn8c:1


Start: 1067| End:1094| Text:'get your Notion database ID' 
Sources:
1. search_internet_cwabyfc5mn8c:1


Start: 1099| End:1137| Text:'add your integration to your database.' 
Sources:
1. search_internet_cwabyfc5mn8c:1


Start: 1223| End:1229| Text:'guide.' 
Sources:
1. search_internet_cwabyfc5mn8c:3
```

Let's ask the agent a final question, this time about tutorials that are relevant for enterprises.

Again, the agent uses the context of the query to decide on the most relevant tool. In this case, it selects the `search_code_examples` tool and provides a response based on the information found.

```python PYTHON
messages = run_agent(
    "Any tutorials that are relevant for enterprises?"
)
```

```mdx

QUESTION:
Any tutorials that are relevant for enterprises?
==================================================
TOOL PLAN:
I will search for 'enterprise tutorials' in the code examples and tutorials tool. 

TOOL CALLS:
Tool name: search_code_examples | Parameters: {"query":"enterprise tutorials"}
==================================================
RESPONSE:
I found a tutorial called 'Advanced Document Parsing For Enterprises'.
==================================================
CITATIONS:

Start: 26| End:69| Text:''Advanced Document Parsing For Enterprises'' 
Sources:
1. search_code_examples_jhh40p32wxpw:4
```

## Summary

In this tutorial, we learned about:

* How to set up tools in an agentic RAG system
* How to run an agentic RAG workflow
* How to automatically route queries to the most relevant data sources

However, so far we have only seen rather simple queries. In practice, we may run into a complex query that needs to simplified, optimized, or split (etc.) before we can perform the retrieval.

In Part 2, we'll learn how to build an agentic RAG system that can expand user queries into parallel queries.