In this chapter, you’ll learn how to build a RAG-powered chatbot that leverages text embeddings using the Chat, Embed, and Rerank endpoints.

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

In the previous chapter, you learned how to get started with RAG using the Chat endpoint.

In this chapter, you’ll learn how to add RAG applications that leverage text embeddings and how to build a chatbot using the V2 Chat API together with Embed and Rerank.

In the previous chapter, we used a short list of simple documents. In real-world applications, however, developers typically need to work with a larger volume of documents that each vary in length.

This is where text embeddings can help. With embeddings, we can split these documents into smaller chunks and build a semantic search system that can retrieve the most relevant chunks to a user query based on contextual meaning, and not just keyword-matching.

Step-by-Step Guide

In V2, RAG is best understood as a workflow rather than a set of endpoint modes:

Grounded generation with documents: Pass relevant documents to the Chat endpoint so the model can generate a grounded response with citations.
Retrieval with Embed and Rerank: Retrieve and rerank relevant documents in your application before calling Chat.
Search query generation: Use application logic or a V2 tool call to turn a user message into one or more retrieval queries when needed.

In this chapter, you’ll build that workflow end to end: generate retrieval queries, retrieve and rerank document chunks, then pass the selected documents to Chat for grounded response generation. Refer to the RAG documentation for more details.

What We'll Build

We’ll build a chatbot that answers users’ questions about the contents in Cohere's documentation on prompt engineering. Let’s examine the demo application's high-level implementation plan (see the diagram below).

An overview of what we'll build

The steps to building a RAG-powered chatbot are summarized below:

Setup phase:

Step 0: Ingest the documents – get documents, chunk, embed, and index

For each user-chatbot interaction:

Step 1: Get the user message
Step 2: Generate one or more retrieval queries using a V2 tool call
If at least one query is generated:
- Step 3: Retrieve and rerank relevant documents
- Step 4: Call the Chat endpoint with the retrieved documents to generate a grounded response with citations
If no query is generated:
- Step 4: Call the Chat endpoint without documents to generate a direct response

Throughout the conversation:

Append the user-chatbot interaction to the conversation thread
Repeat with every interaction

To build a RAG system that can effectively handle a complex corpus of documents, we’ll need to use several Cohere API endpoints, including:

Chat: For handling the main logic of the chatbot, including turning a user message into queries, generating responses, and producing citations
Embed: For turning textual documents into their embeddings representation, later to be used in retrieval (we’ll use the latest Embed v4 model)
Rerank: For reranking the retrieved documents according to their relevance to a query (we’ll use a Rerank v4 model)

This demo application will use Cohere’s Chat, Embed, and Rerank endpoints

For further reading, the API reference page contains a detailed description of the Chat endpoint’s input parameters and response objects.

Setup

First, let’s import the necessary libraries for this project. This includes cohere, hnswlib for the vector library, and unstructured for chunking the documents (more details on these later).

pip install cohere hnswlib unstructured

Then, import the necessary modules from these libraries in addition to other required modules. Let’s also create a Cohere client.

import json
import cohere
import uuid
import hnswlib
from typing import List, Dict
from unstructured.partition.html import partition_html
from unstructured.chunking.title import chunk_by_title

co = cohere.ClientV2("COHERE_API_KEY")

Create the Vectorstore Component

The Vectorstore class handles the ingestion of documents into embeddings (or vectors) and the retrieval of relevant documents given a query.

The Vectorstore component handles document ingestion and retrieval

As an example, we’ll use the contents from Cohere's documentation on prompt engineering. It consists of four web pages, each in the Python list raw_documents below. Each entry is identified by its title and URL.

raw_documents = [
    {
        "title": "Crafting Effective Prompts",
        "url": "https://docs.cohere.com/docs/crafting-effective-prompts"},
    {
        "title": "Advanced Prompt Engineering Techniques",
        "url": "https://docs.cohere.com/docs/advanced-prompt-engineering-techniques"},
    {
        "title": "Prompt Truncation",
        "url": "https://docs.cohere.com/docs/prompt-truncation"},
    {
        "title": "Preambles",
        "url": "https://docs.cohere.com/docs/preambles"}
]

We implement this in the Vectorstore class below, which takes the raw_documents list as input.

class Vectorstore:
    def __init__(self, raw_documents: List[Dict[str, str]]):
        self.raw_documents = raw_documents
        self.docs = []
        self.docs_embs = []
        self.retrieve_top_k = 10
        self.rerank_top_k = 3
        self.load_and_chunk()
        self.embed()
        self.index()

We also initialize a few instance attributes and methods. The attributes include self.raw_documents to represent the raw documents, self.docs to represent the chunked version of the documents, self.docs_embs to represent the embeddings of the chunked documents, and a couple of top_k parameters to be used for retrieval and reranking.

Meanwhile, the methods include load_and_chunk(), embed(), and index() for ingesting raw documents. As you’ll see, we will also specify a retrieve() method to retrieve relevant document chunks given a query.

The document ingestion portion of the Documents component

Load and Chunk the Documents

The load_and_chunk() method loads the raw documents from the URL and breaks them into smaller chunks. Chunking for information retrieval is a broad topic in and of itself, with many strategies being discussed within the AI community. For our example, we’ll utilize the partition_html method from the unstructured library. Read its documentation for more information about its chunking approach.

Each chunk is turned into a dictionary with three fields:

title: The web page’s title
text: The textual content of the chunk
url: The web page’s URL

This information will eventually be passed to the chatbot’s prompt for generating the response, so it’s crucial to populate relevant information into this dictionary. Note that we are not limited to these three fields. At a minimum, the Chat endpoint requires the text field, but beyond that, we can add custom fields that can provide more context about the document, such as subtitles, snippets, tags, and others.

The resulting dictionaries are stored in the self.docs attribute.

class Vectorstore:
    
    ...
    ...    

def load_and_chunk(self) -> None:
        """
        Loads the text from the sources and chunks the HTML content.
        """
        print("Loading documents...")

        for raw_document in self.raw_documents:
            elements = partition_html(url=raw_document["url"])
            chunks = chunk_by_title(elements)
            for chunk in chunks:
                self.docs.append(
                    {
                        "title": raw_document["title"],
                        "text": str(chunk),
                        "url": raw_document["url"],
                    }
                )

Embed the Document Chunks

The embed() method generates embeddings of the chunked documents. We use the Embed endpoint and Cohere's embed-v4.0 model. Since the endpoint has a limit of 96 documents per call, we send them in batches.

With the Embed v4 model, we define an input_type and request float embeddings with embedding_types=["float"]. Since our document chunks will be used for retrieval, we use search_document as the input_type.

The resulting chunk embeddings are stored in the self.docs_embs attribute.

class Vectorstore:
    
    ...
    ...

    def embed(self) -> None:
        """
        Embeds the document chunks using the Cohere API.
        """
        print("Embedding document chunks...")

        batch_size = 90
        self.docs_len = len(self.docs)
        for i in range(0, self.docs_len, batch_size):
            batch = self.docs[i : min(i + batch_size, self.docs_len)]
            texts = [item["text"] for item in batch]
            docs_embs_batch = co.embed(
                model="embed-v4.0",
                input_type="search_document",
                texts=texts,
                embedding_types=["float"],
            ).embeddings.float
            self.docs_embs.extend(docs_embs_batch)

Index Document Chunks

The index() method indexes the document chunk embeddings. We build an index to store the embeddings in a structured and organized way in order to ensure efficient similarity search during retrieval.

There are many options available for building an index. For production environments, typically a vector database (like Weaviate or MongoDB) is required to handle the continuous process of indexing documents and maintaining the index.

In our example, however, we’ll keep it simple and use a vector library instead. We can choose from many open-source projects, such as Faiss, Annoy, ScaNN, or Hnswlib, which is the one we’ll use. These libraries store embeddings in in-memory indexes and implement approximate nearest neighbor (ANN) algorithms to make similarity search efficient.

The resulting document chunk embeddings are stored in the self.idx attribute.

class Vectorstore:
    
    ...
    ...

    def index(self) -> None:
        """
        Indexes the documents for efficient retrieval.
        """
        print("Indexing documents...")

        self.idx = hnswlib.Index(space="ip", dim=1536)
        self.idx.init_index(max_elements=self.docs_len, ef_construction=512, M=64)
        self.idx.add_items(self.docs_embs, list(range(len(self.docs_embs))))

        print(f"Indexing complete with {self.idx.get_current_count()} documents.")

Implement Retrieval

The retrieve() method uses semantic search to retrieve relevant document chunks given a query, and it has two steps: (1) dense retrieval, (2) reranking.

A more detailed view of document ingestion, retrieval, and reranking

Dense Retrieval

We implement a dense retrieval system that leverages embeddings to retrieve document chunks, offering significant improvements over basic keyword-matching approaches. Embeddings can capture the contextual meaning of a document, thus enabling the retrieval of highly relevant results to the given query.

We embed the query using the same embed-v4.0 model that we used to embed the document chunks, but this time, we set input_type=”search_query”.

Search is performed by the knn_query() method from the hnswlib library. Given a query, it returns the document chunks most similar to the query. We define the number of document chunks to return using the attribute self.retrieve_top_k=10.

Reranking

After dense retrieval, we implement a reranking step. While our dense retrieval component is already highly capable of retrieving relevant sources, Cohere Rerank provides an additional boost to the quality of the search results, especially for complex and domain-specific queries. It takes the search results and sorts them according to their relevance to the query.

We call the Rerank endpoint with co.rerank() and pass the query and the list of document chunks to be reranked. We also define the number of top reranked document chunks to retrieve using the attribute self.rerank_top_k=3. The model we use is rerank-v4.0-pro.

This method returns the top retrieved document chunks as a Python list docs_retrieved, so that they can be passed to the chatbot, which we’ll implement next.

class Vectorstore:

    ...
    ...

    def retrieve(self, query: str) -> List[Dict[str, str]]:
        """
        Retrieves document chunks based on the given query.

        Parameters:
        query (str): The query to retrieve document chunks for.

        Returns:
        List[Dict[str, str]]: A list of dictionaries representing the retrieved document chunks, with 'title', 'text', and 'url' keys.
        """

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

        doc_ids = self.idx.knn_query(query_emb, k=self.retrieve_top_k)[0][0]

        # Reranking
        docs_to_rerank = [self.docs[doc_id] for doc_id in doc_ids]

        rerank_results = co.rerank(
            model="rerank-v4.0-pro",
            query=query,
            documents=[f"{doc['title']}\n\n{doc['text']}" for doc in docs_to_rerank],
            top_n=self.rerank_top_k,
        )
        
        doc_ids_reranked = [doc_ids[result.index] for result in rerank_results.results]

        docs_retrieved = []
        for doc_id in doc_ids_reranked:
            docs_retrieved.append(
                {
                    "title": self.docs[doc_id]["title"],
                    "text": self.docs[doc_id]["text"],
                    "url": self.docs[doc_id]["url"],
                }
            )

        return docs_retrieved

Process the Documents

We can now process the raw documents. We do that by creating an instance of Vectorstore. In our case, we get a total of 136 documents, chunked from the four web URLs.

vectorstore = Vectorstore(raw_documents)

Loading documents...
Embedding documents...
Indexing documents...
Indexing complete with 134 documents.

Test Retrieval

Before going further, we first test the document retrieval part of the system. First, we create an instance of the Vectorstore with the raw documents that we have defined. Then, we use the retrieve method to retrieve the most relevant documents to the query "Prompting by giving examples."

vectorstore.retrieve("Prompting by giving examples")

And here’s the response. We can see that the document chunks returned are indeed highly relevant to the query we sent.

[{'title': 'Advanced Prompt Engineering Techniques',
  'text': 'Few-shot Prompting\n\nUnlike the zero-shot examples above, few-shot prompting is a technique that provides a model with examples of the task being performed before asking the specific question to be answered. We can steer the LLM toward a high-quality solution by providing a few relevant and diverse examples in the prompt. Good examples condition the model to the expected response type and style.',
  'url': 'https://docs.cohere.com/docs/advanced-prompt-engineering-techniques'},
 {'title': 'Crafting Effective Prompts',
  'text': 'Incorporating Example Outputs\n\nLLMs respond well when they have specific examples to work from. For example, instead of asking for the salient points of the text and using bullet points “where appropriate”, give an example of what the output should look like.',
  'url': 'https://docs.cohere.com/docs/crafting-effective-prompts'},
 {'title': 'Advanced Prompt Engineering Techniques',
  'text': 'In addition to giving correct examples, including negative examples with a clear indication of why they are wrong can help the LLM learn to distinguish between correct and incorrect responses. Ordering the examples can also be important; if there are patterns that could be picked up on that are not relevant to the correctness of the question, the model may incorrectly pick up on those instead of the semantics of the question itself.',
  'url': 'https://docs.cohere.com/docs/advanced-prompt-engineering-techniques'}]

Run the Chatbot

We can now run the chatbot. For this, we create a run_chatbot function which includes the RAG components:

Search query generation: For each user message, we call the V2 Chat endpoint with a query-generation tool that returns one or more retrieval queries, or an empty list when retrieval is not needed.
Document retrieval: If there is no search query generated, we call the Chat endpoint to generate a response directly. If there is at least one, we call the retrieve method from the Vectorstore instance to retrieve the most relevant documents to each query.
Response and citation generation: Finally, all the results from all queries are appended to a list and passed to the Chat endpoint for response generation. We print the streamed response, together with the citations, for easy reference.

The Chatbot component handles the chatbot logic, from getting the user message to generating the response

The generate_chat function handles the chatbot logic, from getting the user message to generating the response

QUERY_MODEL = "command-a-03-2025"
CHAT_MODEL = "command-a-plus-05-2026"

query_generation_tool = [
    {
        "type": "function",
        "function": {
            "name": "generate_search_queries",
            "description": "Return search queries for retrieval, or an empty list if retrieval is not needed.",
            "parameters": {
                "type": "object",
                "properties": {
                    "queries": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Search queries. Empty if retrieval is not needed.",
                    }
                },
                "required": ["queries"],
            },
        },
    }
]


def generate_search_queries(message: str, chat_history: List[Dict[str, str]]) -> List[str]:
    instructions = (
        "Return search queries that will find helpful information for answering the user question accurately. "
        "Return an empty queries list if retrieval is not needed."
    )
    response = co.chat(
        model=QUERY_MODEL,
        messages=[
            {"role": "system", "content": instructions},
            *chat_history,
            {"role": "user", "content": message},
        ],
        tools=query_generation_tool,
        tool_choice="REQUIRED",
        temperature=0.1,
    )

    search_queries = []
    for tool_call in response.message.tool_calls or []:
        arguments = json.loads(tool_call.function.arguments)
        search_queries.extend(arguments.get("queries", []))
    return search_queries


def format_documents_for_chat(documents: List[Dict[str, str]]) -> List[Dict[str, Dict[str, str]]]:
    return [
        {
            "id": f"doc_{idx}",
            "data": {
                "title": document["title"],
                "text": document["text"],
                "url": document["url"],
            },
        }
        for idx, document in enumerate(documents)
    ]


def run_chatbot(message: str, chat_history: List[Dict[str, str]] | None = None):
    if chat_history is None:
        chat_history = []

    messages = [*chat_history, {"role": "user", "content": message}]
    search_queries = generate_search_queries(message, chat_history)

    documents = []
    if search_queries:
        print("Retrieving information...", end="")
        for query in search_queries:
            documents.extend(vectorstore.retrieve(query))

    chat_request = {
        "model": CHAT_MODEL,
        "messages": [
            {
                "role": "system",
                "content": "Answer concisely. When documents are provided, ground the response in those documents and include citations.",
            },
            *messages,
        ],
    }
    if documents:
        chat_request["documents"] = format_documents_for_chat(documents)

    response = co.chat_stream(**chat_request)

    chatbot_response = ""
    citations = []
    print("\nChatbot:")

    for event in response:
        if not event:
            continue
        if event.type == "content-delta":
            text = event.delta.message.content.text
            if text:
                print(text, end="")
                chatbot_response += text
        if event.type == "citation-start":
            citations.append(event.delta.message.citations)

    if citations:
        print("\n\nCITATIONS:")
        for citation in citations:
            print(citation)

    chat_history = messages + [{"role": "assistant", "content": chatbot_response}]
    return chat_history

Search Query Generation

Let's take a deeper look at the search query generation feature. Based on the user message, the chatbot needs to decide if it needs to consult external information before responding. If so, the chatbot determines an optimal set of search queries to use for retrieval. In V2, this is implemented with a query-generation tool call that returns a list of search queries, or an empty list when retrieval is not needed.

The generated queries can be accessed from the tool call arguments. To understand how this works, let’s look at a few scenarios:

No query needed: Suppose we have a user message of “Hello, I need help with a report I'm writing”. This type of message doesn’t require any additional context from external information, so retrieval is not required. A direct chatbot response will suffice (for example: “Sure, how can I help?”). When we send this to the Chat endpoint, we get an empty search_queries result, which is what we expect.
One query generated: Take this user message: "What did the report say about the company's Q4 performance?” This does require additional context as it refers to a report, hence retrieval is required. Given this message, the Chat endpoint returns the search_queries result of Q4 company performance. Here it turns the user message into a query optimized for search. Another important scenario is generating queries in the context of the conversation. Suppose there’s an ongoing conversation where the user is learning from the chatbot about deep learning. If at some point, the user asks, “Why is it important”, then the generated search_queries will become why is deep learning important, providing the much-needed context for the retrieval process.
More than one query generated: What if the user message is a bit more complex, such as "What did the report say about the company's Q4 performance and its range of products and services?” This requires multiple pieces of information to be retrieved. Given this message, the Chat endpoint returns two search_queries results: Q4 company performance and company's range of products and services.

These scenarios highlight the adaptability of the Chat endpoint to decide on the next course of action based on a user message.

Document Retrieval

Let's take a deeper look at the document retrieval step. What happens next depends on how many search queries are returned.

If search queries are returned

If the chatbot response contains at least one search query, we call the retrieve() method from the Vectorstore class instance to retrieve document chunks that are relevant to the queries.

Then, we call the Chat endpoint to generate a response, adding a documents parameter to the call to pass the relevant document chunks.

If no search queries are returned

Meanwhile, if the chatbot response doesn’t contain any search queries, then it doesn’t require information retrieval. To generate the response, we call the Chat endpoint another time, passing the user message and without needing to add any sources to the call.

In either case, we pass the accumulated messages list, which retains the interactions between the user and the chatbot in the same conversation thread. We also use the chat_stream endpoint so we can stream the chatbot response to the application.

Response and Citation Generation

Let's take a deeper look at the response generation step. The chatbot response includes a stream of events, such as the generated text and citation events.

To display the response, we use the content-delta events from the response stream.

The citation-start events indicate the spans of text from the retrieved document chunks on which the response is grounded. Here is one example:

Citation(start=382, end=397, text='similar vectors', sources=[DocumentSource(type='document', id='doc_0', document={...})])

The format of each citation is:

start: The starting point of a span where one or more documents are referenced
end: The ending point of a span where one or more documents are referenced
text: The text representing this span
sources: The document chunks being referenced, including their IDs and document metadata

Example conversation

Here’s an example of a conversation that happens over a few turns:

# Turn # 1
chat_history = run_chatbot("Hello, I have a question")

Chatbot:
Of course! I am here to help. Please go ahead with your question, and I will do my best to assist you.

# Turn # 2
chat_history = run_chatbot("What's the difference between zero-shot and few-shot prompting", chat_history)

Retrieving information...
Chatbot:
Zero-shot prompting involves asking the model to perform a task without providing any examples. On the other hand, few-shot prompting is a technique where the model is provided with a few relevant and diverse examples of the task being performed before asking the specific question to be answered. These examples help steer the model toward a high-quality solution and condition it to the expected response type and style.

CITATIONS:
Citation(start=0, end=19, text='Zero-shot prompting', sources=[DocumentSource(type='document', id='doc_0', document={...})])
Citation(start=115, end=133, text='few-shot prompting', sources=[DocumentSource(type='document', id='doc_0', document={...})])

# Turn # 3
chat_history = run_chatbot("How would the latter help?", chat_history)

Retrieving information...
Chatbot:
Few-shot prompting can vastly improve the quality of the model's completions. Providing a few relevant and diverse examples helps steer the model toward a high-quality solution by conditioning it to the expected response type and style.

CITATIONS:
Citation(start=23, end=77, text="vastly improve the quality of the model's completions.", sources=[DocumentSource(type='document', id='doc_2', document={...})])
Citation(start=90, end=123, text='few relevant and diverse examples', sources=[DocumentSource(type='document', id='doc_0', document={...})])

# Turn # 4
chat_history = run_chatbot("What do you know about 5G networks?", chat_history)

Retrieving information...
Chatbot:
Sorry, I don't have any information about 5G networks. Can I help you with anything else?
There are a few observations worth pointing out:

Direct response: For user messages that don’t require retrieval (“Hello, I have a question”), the chatbot responds directly without requiring retrieval.
Citation generation: For responses that do require retrieval ("What's the difference between zero-shot and few-shot prompting"), the endpoint returns the response together with the citations. These are fine-grained citations, which means they refer to specific spans of the generated text.
State management: The application maintains the state of the conversation via the `chat_history` messages list, for example, by correctly responding to a vague user message such as "How would the latter help?"
Response synthesis: The model can decide if none of the retrieved documents provide the necessary information to answer a user message. For example, when asked the question, “What do you know about 5G networks”, the chatbot retrieves external information from the index. However, it doesn’t use any of the information in its response as none of it is relevant to the question.

In the conversation above, notice a few observations that reflect the different components of what we built:

Direct response: For user messages that don’t require retrieval (“Hello, I have a question”), the chatbot responds directly without requiring retrieval.
Citation generation: For responses that do require retrieval (“What’s the difference between word and sentence embeddings”), the endpoint can stream the response together with fine-grained citations.
State management: The application maintains the state of the conversation via the chat_history messages list, for example, by being able to correctly respond to a vague user message of “And what are their similarities”
Response synthesis: The model can decide if none of the retrieved documents provide the necessary information required to answer a user message. For example, when asked the question, “What do you know about 5G networks”, the chatbot goes on and retrieves external information from the index. However, it doesn’t use any of the information in its response as none of them is relevant to the question.

Here are the contents of the chat history.

print("Chat history:")
for c in chat_history:
    print(c, "\n")

Chat history:
{'role': 'user', 'content': 'Hello, I have a question'}

{'role': 'assistant', 'content': 'Of course! I am here to help. Please go ahead with your question, and I will do my best to assist you.'}

{'role': 'user', 'content': "What's the difference between zero-shot and few-shot prompting"}

{'role': 'assistant', 'content': 'Zero-shot prompting involves asking the model to perform a task without providing any examples...'}

{'role': 'user', 'content': 'How would the latter help?'}

{'role': 'assistant', 'content': "Few-shot prompting can vastly improve the quality of the model's completions..."}

{'role': 'user', 'content': 'What do you know about 5G networks?'}

{'role': 'assistant', 'content': "Sorry, I don't have any information about 5G networks. Can I help you with anything else?"}

Conclusion

In this chapter, you learned how to build a RAG-powered chatbot with the Chat endpoint. With access to a collection of documents, the chatbot is able to provide contextually relevant responses to user requests, along with verifiable citations.

Contents

How to Build a RAG-Powered Chatbot with Chat, Embed, and Rerank