Overview

RAG systems often struggle with retrieving truly relevant documents, leading to poor response quality. Purely embedding-based retrieval may miss nuanced connections between queries and documents. This chapter shows how to combine Cohere's Chat, Embed, and Rerank APIs to build a more effective agentic RAG pipeline.

This chapter demonstrates how to create an agentic RAG system that leverages the tool use capabilities of the Chat API. It generates optimized search queries based on a user query, retrieves documents through embeddings, reranks results for relevance, and generates grounded responses.

By the end, you will have built an agentic RAG pipeline that significantly improves document retrieval accuracy and response quality.

What we'll cover

Creating a document search tool
Implementing embed and rerank pipeline
Handling user queries with tool calls
Processing search results and reranking
Generating responses with automatic citations

import cohere
import os
co = cohere.ClientV2(os.getenv("COHERE_API_KEY"))

import json
import numpy as np

Creating a document search tool

First, we'll define a function that will search for documents in our corpus. We start by definining a small set of documents that the tool will have access to, which is a list of company FAQ documents that might be useful for a new employee.

company_faqs = [
    {
        "data": {
            "text": "Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged."
        }
    },
    {
        "data": {
            "text": "Finding Coffee Spots: For your caffeine fix, head to the break room's coffee machine or cross the street to the café for artisan coffee."
        }
    },
    {
        "data": {
            "text": "Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!"
        }
    },
    {
        "data": {
            "text": "Working Hours Flexibility: We prioritize work-life balance. While our core hours are 9 AM to 5 PM, we offer flexibility to adjust as needed."
        }
    },
    {
        "data": {
            "text": "Side Projects Policy: We encourage you to pursue your passions. Just be mindful of any potential conflicts of interest with our business."
        }
    },
    {
        "data": {
            "text": "Reimbursing Travel Expenses: Easily manage your travel expenses by submitting them through our finance tool. Approvals are prompt and straightforward."
        }
    },
    {
        "data": {
            "text": "Working from Abroad: Working remotely from another country is possible. Simply coordinate with your manager and ensure your availability during core hours."
        }
    },
    {
        "data": {
            "text": "Health and Wellness Benefits: We care about your well-being and offer gym memberships, on-site yoga classes, and comprehensive health insurance."
        }
    },
    {
        "data": {
            "text": "Performance Reviews Frequency: We conduct informal check-ins every quarter and formal performance reviews twice a year."
        }
    },
    {
        "data": {
            "text": "Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead."
        }
    },
]

Next, we create an text embedding function to embed the documents using the Embed endpoint.

def embed_texts(documents):
    """Embed a list of texts using Cohere's embed model."""

    embeddings = co.embed(
        texts=documents,
        model="embed-v4.0",
        input_type="search_document",
        embedding_types=["float"],
    ).embeddings.float
    return embeddings

doc_emb = embed_texts([doc["data"]["text"] for doc in company_faqs])

Now we can implement a complete document search function that does the following:

Takes a search query and embeds it
Retrieves relevant documents through embeddings
Reranks the results for relevance
Returns the top n reranked documents

def document_search(query):

    def retrieve_documents(query, documents, doc_emb):

        query_emb = co.embed(
            model="embed-v4.0",
            input_type="search_query",
            texts=[query],
            embedding_types=["float"],
        ).embeddings.float

        # Compute dot product similarity and display results
        n = 5
        scores = np.dot(query_emb, np.transpose(doc_emb))[0]
        max_idx = np.argsort(-scores)[:n]

        retrieved_documents = [documents[item] for item in max_idx]

        print("=" * 100)
        print(f"QUERY: {query}")
        print("=" * 100)
        print("First stage retrieval")
        for rank, idx in enumerate(max_idx):
            doc = retrieved_documents[rank]
            print(f"Rank: {rank+1}")
            print(f"Score: {scores[idx]}")
            print(f"Document: {doc}")
        print("-" * 80)

        return retrieved_documents

    def rerank_documents(query, retrieved_documents):
        # Rerank the documents
        results = co.rerank(
            model="rerank-v3.5",
            query=query,
            documents=[doc["data"]["text"] for doc in retrieved_documents],
            top_n=2,
        )

        # Display the reranking results
        print("Second stage reranking")
        for idx, result in enumerate(results.results):
            print(f"Rank: {idx+1}")
            print(f"Score: {result.relevance_score}")
            print(f"Document: {retrieved_documents[result.index]}\n")

        # Return reranked documents in the same format as retrieved_documents
        reranked_documents = []
        for result in results.results:
            doc = retrieved_documents[result.index]
            reranked_documents.append(doc)

        return reranked_documents

    documents = retrieve_documents(query, company_faqs, doc_emb)
    documents_reranked = rerank_documents(query, company_faqs)

    return documents_reranked


functions_map = {
    "document_search": document_search,
}

We then create the tool definition to represent this document search function. Here, the tool contains one function called document_search that generates one or more queries based on the user query, sends the queries to the document_search function, and gets back the top n reranked documents.

document_search_tool = {
    "type": "function",
    "function": {
        "name": "document_search",
        "description": "Returns a list of relevant documents for a given text query.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "A query to search the document collection with.",
                }
            },
            "required": ["query"],
        },
    },
}

tools = [document_search_tool]

1 - User query

The first part of the workflow is to get the user query.

user_message = "How to get to know my teammates"

messages = [
    {"role": "user", "content": user_message},
]

2 - Tool calling

Next, we perform tool calling to the document_search tool, if required, based on the user query.

response = co.chat(
    model="command-a-03-2025", messages=messages, tools=tools
)

if response.message.tool_calls:
    messages.append(
        {
            "role": "assistant",
            "tool_plan": response.message.tool_plan,
            "tool_calls": response.message.tool_calls,
        }
    )
    print(response.message.tool_plan, "\n")
    print(response.message.tool_calls)

I will search for how to get to know my teammates and where to get coffee. 

[ToolCallV2(id='document_search_m1q6a0ngb2mf', type='function', function=ToolCallV2Function(name='document_search', arguments='{"query":"how to get to know my teammates"}')), ToolCallV2(id='document_search_b4fhmj148k83', type='function', function=ToolCallV2Function(name='document_search', arguments='{"query":"where to get coffee"}'))]

3 - Tool results

The retrieved results are then passed back and added to the messages list, which will then be passed to the Chat endpoint to generate the response.

# Calling the tool

if response.message.tool_calls:
    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:
            # Optional: the "document" object can take an "id" field for use in citations, otherwise auto-generated
            tool_content.append(
                {
                    "type": "document",
                    "document": {"data": json.dumps(data)},
                }
            )
        messages.append(
            {
                "role": "tool",
                "tool_call_id": tc.id,
                "content": tool_content,
            }
        )

====================================================================================================
QUERY: how to get to know my teammates
====================================================================================================
First stage retrieval
Rank: 1
Score: 0.3264442589056174
Document: {'data': {'text': 'Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!'}}
Rank: 2
Score: 0.2912684030497781
Document: {'data': {'text': 'Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead.'}}
Rank: 3
Score: 0.2062645307053807
Document: {'data': {'text': 'Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.'}}
Rank: 4
Score: 0.1674031603572354
Document: {'data': {'text': "Finding Coffee Spots: For your caffeine fix, head to the break room's coffee machine or cross the street to the café for artisan coffee."}}
Rank: 5
Score: 0.15291999832696462
Document: {'data': {'text': 'Performance Reviews Frequency: We conduct informal check-ins every quarter and formal performance reviews twice a year.'}}
--------------------------------------------------------------------------------
Second stage reranking
Rank: 1
Score: 0.064212106
Document: {'data': {'text': 'Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.'}}

Rank: 2
Score: 0.061317287
Document: {'data': {'text': 'Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!'}}

====================================================================================================
QUERY: where to get coffee
====================================================================================================
First stage retrieval
Rank: 1
Score: 0.43517191609702893
Document: {'data': {'text': "Finding Coffee Spots: For your caffeine fix, head to the break room's coffee machine or cross the street to the café for artisan coffee."}}
Rank: 2
Score: 0.1496811430092236
Document: {'data': {'text': 'Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.'}}
Rank: 3
Score: 0.14092252045479745
Document: {'data': {'text': 'Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead.'}}
Rank: 4
Score: 0.10707084742007339
Document: {'data': {'text': 'Health and Wellness Benefits: We care about your well-being and offer gym memberships, on-site yoga classes, and comprehensive health insurance.'}}
Rank: 5
Score: 0.09660644742314248
Document: {'data': {'text': 'Working Hours Flexibility: We prioritize work-life balance. While our core hours are 9 AM to 5 PM, we offer flexibility to adjust as needed.'}}
--------------------------------------------------------------------------------
Second stage reranking
Rank: 1
Score: 0.6213807
Document: {'data': {'text': "Finding Coffee Spots: For your caffeine fix, head to the break room's coffee machine or cross the street to the café for artisan coffee."}}

Rank: 2
Score: 0.021617446
Document: {'data': {'text': 'Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead.'}}

4 - Response and citations

The Chat endpoint then uses the documents retrieved from the document_search tool to generate a response and citations.

response = co.chat(
    model="command-a-03-2025", messages=messages, tools=tools
)

messages.append(
    {"role": "assistant", "content": response.message.content[0].text}
)

print("RESPONSE")
print(response.message.content[0].text)
print("-"*60)

print("CITATIONS")
print(response.message.citations)

RESPONSE
To get to know your teammates, you can join relevant Slack channels to stay informed and engaged. You can also participate in team-building activities such as monthly outings and weekly game nights.

For coffee, you can head to the break room's coffee machine or cross the street to the café for artisan coffee.
------------------------------------------------------------
CITATIONS
[Citation(start=39, end=67, text='join relevant Slack channels', sources=[ToolSource(type='tool', id='document_search_m1q6a0ngb2mf:0', tool_output={'data': '{"text":"Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged."}'})], type='TEXT_CONTENT'), Citation(start=71, end=97, text='stay informed and engaged.', sources=[ToolSource(type='tool', id='document_search_m1q6a0ngb2mf:0', tool_output={'data': '{"text":"Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged."}'})], type='TEXT_CONTENT'), Citation(start=111, end=150, text='participate in team-building activities', sources=[ToolSource(type='tool', id='document_search_m1q6a0ngb2mf:1', tool_output={'data': '{"text":"Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!"}'})], type='TEXT_CONTENT'), Citation(start=159, end=198, text='monthly outings and weekly game nights.', sources=[ToolSource(type='tool', id='document_search_m1q6a0ngb2mf:1', tool_output={'data': '{"text":"Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!"}'})], type='TEXT_CONTENT'), Citation(start=220, end=259, text="head to the break room's coffee machine", sources=[ToolSource(type='tool', id='document_search_b4fhmj148k83:0', tool_output={'data': '{"text":"Finding Coffee Spots: For your caffeine fix, head to the break room\'s coffee machine or cross the street to the café for artisan coffee."}'})], type='TEXT_CONTENT'), Citation(start=263, end=311, text='cross the street to the café for artisan coffee.', sources=[ToolSource(type='tool', id='document_search_b4fhmj148k83:0', tool_output={'data': '{"text":"Finding Coffee Spots: For your caffeine fix, head to the break room\'s coffee machine or cross the street to the café for artisan coffee."}'})], type='TEXT_CONTENT')]

for message in messages:
    print(message, "\n")

{'role': 'user', 'content': 'How to get to know my teammates and where to get coffee'} 

{'role': 'assistant', 'tool_plan': 'I will search for how to get to know my teammates and where to get coffee.', 'tool_calls': [ToolCallV2(id='document_search_m1q6a0ngb2mf', type='function', function=ToolCallV2Function(name='document_search', arguments='{"query":"how to get to know my teammates"}')), ToolCallV2(id='document_search_b4fhmj148k83', type='function', function=ToolCallV2Function(name='document_search', arguments='{"query":"where to get coffee"}'))]} 

{'role': 'tool', 'tool_call_id': 'document_search_m1q6a0ngb2mf', 'content': [{'type': 'document', 'document': {'data': '{"data": {"text": "Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged."}}'}}, {'type': 'document', 'document': {'data': '{"data": {"text": "Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!"}}'}}]} 

{'role': 'tool', 'tool_call_id': 'document_search_b4fhmj148k83', 'content': [{'type': 'document', 'document': {'data': '{"data": {"text": "Finding Coffee Spots: For your caffeine fix, head to the break room\'s coffee machine or cross the street to the caf\\u00e9 for artisan coffee."}}'}}, {'type': 'document', 'document': {'data': '{"data": {"text": "Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead."}}'}}]} 

{'role': 'assistant', 'content': "To get to know your teammates, you can join relevant Slack channels to stay informed and engaged. You can also participate in team-building activities such as monthly outings and weekly game nights.\n\nFor coffee, you can head to the break room's coffee machine or cross the street to the café for artisan coffee."}

Agentic RAG with Chat, Embed, and Rerank