notebooks/llmu/RAG_with_Chat_Embed_and_Rerank.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "ctaLvRUsfpj8",
      "metadata": {
        "id": "ctaLvRUsfpj8"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/cohere-ai/notebooks/blob/main/notebooks/llmu/RAG_with_Chat_Embed_and_Rerank.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "61bac5b5",
      "metadata": {
        "id": "61bac5b5"
      },
      "source": [
        "# RAG with Chat, Embed, and Rerank\n",
        "\n",
        "This notebook shows how to build a RAG-powered chatbot with Cohere's Chat endpoint.  The chatbot can extract relevant information from external documents and produce verifiable, inline citations in its responses.\n",
        "\n",
        "Read the accompanying [article here](https://txt.cohere.com/rag-chatbot/).\n",
        "\n",
        "This application will use several Cohere API endpoints:\n",
        "\n",
        "- Chat: For handling the main logic of the chatbot, including turning a user message into queries, generating responses, and producing citations\n",
        "- Embed: For turning textual documents into their embeddings representation, later to be used in retrieval (we’ll use the latest, state-of-the-art Embed v3 model)\n",
        "- Rerank: For reranking the retrieved documents according to their relevance to a query\n",
        "\n",
        "The diagram below provides an overview of what we’ll build."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "33327522",
      "metadata": {},
      "source": [
        "![Workflow](../images/llmu/rag/rag-workflow-2.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f6ab2d5d",
      "metadata": {
        "id": "f6ab2d5d"
      },
      "source": [
        "Here is a summary of the steps involved.\n",
        "\n",
        "Initial phase:\n",
        "- **Step 0**: Ingest the documents – get documents, chunk, embed, and index.\n",
        "\n",
        "For each user-chatbot interaction:\n",
        "- **Step 1**: Get the user message\n",
        "- **Step 2**: Call the Chat endpoint in query-generation mode\n",
        "- If at least one query is generated\n",
        "  - **Step 3**: Retrieve and rerank relevant documents\n",
        "  - **Step 4**: Call the Chat endpoint in document mode to generate a grounded response with citations\n",
        "- If no query is generated\n",
        "  - **Step 4**: Call the Chat endpoint in normal mode to generate a response"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "TWyo_5WoNUM-",
      "metadata": {
        "id": "TWyo_5WoNUM-"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5pLAhQmTOKiV",
      "metadata": {
        "id": "5pLAhQmTOKiV"
      },
      "outputs": [],
      "source": [
        "! pip install cohere hnswlib unstructured --upgrade nltk -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "f3a03a57",
      "metadata": {
        "id": "f3a03a57"
      },
      "outputs": [],
      "source": [
        "import cohere\n",
        "import uuid\n",
        "import hnswlib\n",
        "from typing import List, Dict\n",
        "from unstructured.partition.html import partition_html\n",
        "from unstructured.chunking.title import chunk_by_title\n",
        "\n",
        "co = cohere.Client(\"COHERE_API_KEY\") # Get your API key here: https://dashboard.cohere.com/api-keys"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "Dx1cncziCWBB",
      "metadata": {
        "cellView": "form",
        "id": "Dx1cncziCWBB"
      },
      "outputs": [],
      "source": [
        "#@title Enable text wrapping in Google Colab\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9d34e4b7",
      "metadata": {},
      "source": [
        "# Create a vector store for ingestion and retrieval\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "588ed6d0",
      "metadata": {},
      "source": [
        "![RAG components - Vectorstore](../images/llmu/rag/rag-components-vectorstore.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2f7e7d1c",
      "metadata": {
        "id": "2f7e7d1c"
      },
      "source": [
        "\n",
        "First, we define the list of documents we want to ingest and make available for retrieval. As an example, we'll use the contents from the first module of Cohere's *LLM University: What are Large Language Models?*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "3dca4a88",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "3dca4a88",
        "outputId": "b05da1ee-0456-4387-c232-a43e0ffed54c"
      },
      "outputs": [],
      "source": [
        "raw_documents = [\n",
        "    {\n",
        "        \"title\": \"Crafting Effective Prompts\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/crafting-effective-prompts\"},\n",
        "    {\n",
        "        \"title\": \"Advanced Prompt Engineering Techniques\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/advanced-prompt-engineering-techniques\"},\n",
        "    {\n",
        "        \"title\": \"Prompt Truncation\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/prompt-truncation\"},\n",
        "    {\n",
        "        \"title\": \"Preambles\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/preambles\"}\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5e2a8968",
      "metadata": {
        "id": "5e2a8968"
      },
      "source": [
        "Usually the number of documents for practical applications is vast, and so we'll need to be able to search documents efficiently.  This involves breaking the documents into chunks, generating embeddings, and indexing the embeddings, as shown in the image below.  \n",
        "\n",
        "We implement this in the `Vectorstore` class below, which takes the `raw_documents` list as input.  Three methods are immediately called when creating an object of the `Vectorstore` class:\n",
        "\n",
        "\n",
        "`load_and_chunk()`  \n",
        "This method uses the `partition_html()` method from the `unstructured` library to load the documents from URL and break them into smaller chunks.  Each chunk is turned into a dictionary object with three fields:\n",
        "- `title` - the web page’s title,\n",
        "- `text` - the textual content of the chunk, and\n",
        "- `url` - the web page’s URL.  \n",
        "  \n",
        "  \n",
        "`embed()`  \n",
        "This method uses Cohere's `embed-english-v3.0` model to generate embeddings of the chunked documents.  Since our documents will be used for retrieval, we set `input_type=\"search_document\"`.  We send the documents to the Embed endpoint in batches, because the endpoint has a limit of 96 documents per call.\n",
        "\n",
        "`index()`  \n",
        "This method uses the `hsnwlib` package to index the document chunk embeddings.  This will ensure efficient similarity search during retrieval.  Note that `hnswlib` uses a vector library, and we have chosen it for its simplicity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "c28df755",
      "metadata": {},
      "outputs": [],
      "source": [
        "class Vectorstore:\n",
        "    \"\"\"\n",
        "    A class representing a collection of documents indexed into a vectorstore.\n",
        "\n",
        "    Parameters:\n",
        "    raw_documents (list): A list of dictionaries representing the sources of the raw documents. Each dictionary should have 'title' and 'url' keys.\n",
        "\n",
        "    Attributes:\n",
        "    raw_documents (list): A list of dictionaries representing the raw documents.\n",
        "    docs (list): A list of dictionaries representing the chunked documents, with 'title', 'text', and 'url' keys.\n",
        "    docs_embs (list): A list of the associated embeddings for the document chunks.\n",
        "    docs_len (int): The number of document chunks in the collection.\n",
        "    idx (hnswlib.Index): The index used for document retrieval.\n",
        "\n",
        "    Methods:\n",
        "    load_and_chunk(): Loads the data from the sources and partitions the HTML content into chunks.\n",
        "    embed(): Embeds the document chunks using the Cohere API.\n",
        "    index(): Indexes the document chunks for efficient retrieval.\n",
        "    retrieve(): Retrieves document chunks based on the given query.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, raw_documents: List[Dict[str, str]]):\n",
        "        self.raw_documents = raw_documents\n",
        "        self.docs = []\n",
        "        self.docs_embs = []\n",
        "        self.retrieve_top_k = 10\n",
        "        self.rerank_top_k = 3\n",
        "        self.load_and_chunk()\n",
        "        self.embed()\n",
        "        self.index()\n",
        "\n",
        "\n",
        "    def load_and_chunk(self) -> None:\n",
        "        \"\"\"\n",
        "        Loads the text from the sources and chunks the HTML content.\n",
        "        \"\"\"\n",
        "        print(\"Loading documents...\")\n",
        "\n",
        "        for raw_document in self.raw_documents:\n",
        "            elements = partition_html(url=raw_document[\"url\"])\n",
        "            chunks = chunk_by_title(elements)\n",
        "            for chunk in chunks:\n",
        "                self.docs.append(\n",
        "                    {\n",
        "                        \"title\": raw_document[\"title\"],\n",
        "                        \"text\": str(chunk),\n",
        "                        \"url\": raw_document[\"url\"],\n",
        "                    }\n",
        "                )\n",
        "\n",
        "    def embed(self) -> None:\n",
        "        \"\"\"\n",
        "        Embeds the document chunks using the Cohere API.\n",
        "        \"\"\"\n",
        "        print(\"Embedding document chunks...\")\n",
        "\n",
        "        batch_size = 90\n",
        "        self.docs_len = len(self.docs)\n",
        "        for i in range(0, self.docs_len, batch_size):\n",
        "            batch = self.docs[i : min(i + batch_size, self.docs_len)]\n",
        "            texts = [item[\"text\"] for item in batch]\n",
        "            docs_embs_batch = co.embed(\n",
        "                texts=texts, model=\"embed-english-v3.0\", input_type=\"search_document\"\n",
        "            ).embeddings\n",
        "            self.docs_embs.extend(docs_embs_batch)\n",
        "\n",
        "    def index(self) -> None:\n",
        "        \"\"\"\n",
        "        Indexes the document chunks for efficient retrieval.\n",
        "        \"\"\"\n",
        "        print(\"Indexing document chunks...\")\n",
        "\n",
        "        self.idx = hnswlib.Index(space=\"ip\", dim=1024)\n",
        "        self.idx.init_index(max_elements=self.docs_len, ef_construction=512, M=64)\n",
        "        self.idx.add_items(self.docs_embs, list(range(len(self.docs_embs))))\n",
        "\n",
        "        print(f\"Indexing complete with {self.idx.get_current_count()} document chunks.\")\n",
        "\n",
        "    def retrieve(self, query: str) -> List[Dict[str, str]]:\n",
        "        \"\"\"\n",
        "        Retrieves document chunks based on the given query.\n",
        "\n",
        "        Parameters:\n",
        "        query (str): The query to retrieve document chunks for.\n",
        "\n",
        "        Returns:\n",
        "        List[Dict[str, str]]: A list of dictionaries representing the retrieved document chunks, with 'title', 'text', and 'url' keys.\n",
        "        \"\"\"\n",
        "\n",
        "        # Dense retrieval\n",
        "        query_emb = co.embed(\n",
        "            texts=[query], model=\"embed-english-v3.0\", input_type=\"search_query\"\n",
        "        ).embeddings\n",
        "        \n",
        "        doc_ids = self.idx.knn_query(query_emb, k=self.retrieve_top_k)[0][0]\n",
        "\n",
        "        # Reranking\n",
        "        rank_fields = [\"title\", \"text\"] # We'll use the title and text fields for reranking\n",
        "\n",
        "        docs_to_rerank = [self.docs[doc_id] for doc_id in doc_ids]\n",
        "        rerank_results = co.rerank(\n",
        "            query=query,\n",
        "            documents=docs_to_rerank,\n",
        "            top_n=self.rerank_top_k,\n",
        "            model=\"rerank-english-v3.0\",\n",
        "            rank_fields=rank_fields\n",
        "        )\n",
        "\n",
        "        doc_ids_reranked = [doc_ids[result.index] for result in rerank_results.results]\n",
        "\n",
        "        docs_retrieved = []\n",
        "        for doc_id in doc_ids_reranked:\n",
        "            docs_retrieved.append(\n",
        "                {\n",
        "                    \"title\": self.docs[doc_id][\"title\"],\n",
        "                    \"text\": self.docs[doc_id][\"text\"],\n",
        "                    \"url\": self.docs[doc_id][\"url\"],\n",
        "                }\n",
        "            )\n",
        "\n",
        "        return docs_retrieved"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e1bf5d85",
      "metadata": {
        "id": "e1bf5d85"
      },
      "source": [
        "In the code cell below, we initialize an instance of the `Vectorstore` class and pass in the `raw_documents` list as input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "4643e630",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 178
        },
        "id": "4643e630",
        "outputId": "fe01fcb6-3574-4322-d8d0-57d37aad397d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading documents...\n",
            "Embedding document chunks...\n",
            "Indexing document chunks...\n",
            "Indexing complete with 44 document chunks.\n"
          ]
        }
      ],
      "source": [
        "# Create an instance of the Vectorstore class with the given sources\n",
        "vectorstore = Vectorstore(raw_documents)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "61928287",
      "metadata": {
        "id": "61928287"
      },
      "source": [
        "The `Vectorstore` class also has a `retrieve()` method, which we'll use to retrieve relevant document chunks given a query (as in Step 3 in the diagram shared at the beginning of this notebook).  This method has two components: (1) dense retrieval, and (2) reranking.\n",
        "\n",
        "### Dense retrieval\n",
        "\n",
        "First, we embed the query using the same `embed-english-v3.0` model we used to embed the document chunks, but this time we set `input_type=\"search_query\"`.\n",
        "\n",
        "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 can define the number of document chunks to return using the attribute `self.retrieve_top_k=10`.\n",
        "\n",
        "### Reranking\n",
        "\n",
        "After semantic search, we implement a reranking step.  While our semantic search component is already highly capable of retrieving relevant sources, the [Rerank endpoint](https://cohere.com/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.\n",
        "\n",
        "We call the Rerank endpoint with the `co.rerank()` method and define the number of top reranked document chunks to retrieve using the attribute `self.rerank_top_k=3`.  The model we use is `rerank-english-v2.0`.  \n",
        "\n",
        "This method returns the top retrieved document chunks `chunks_retrieved` so that they can be passed to the chatbot.\n",
        "\n",
        "In the code cell below, we check the document chunks that are retrieved for the query `\"multi-head attention definition\"`."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "OwozNf_uPEyX",
      "metadata": {
        "id": "OwozNf_uPEyX"
      },
      "source": [
        "## Test Retrieval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "82617b91",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 267
        },
        "id": "82617b91",
        "outputId": "7f1f2bc8-8ed9-4190-bd6b-7af2d9dc1980"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[{'title': 'Advanced Prompt Engineering Techniques',\n",
              "  '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.',\n",
              "  'url': 'https://docs.cohere.com/docs/advanced-prompt-engineering-techniques'},\n",
              " {'title': 'Crafting Effective Prompts',\n",
              "  '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.',\n",
              "  'url': 'https://docs.cohere.com/docs/crafting-effective-prompts'},\n",
              " {'title': 'Advanced Prompt Engineering Techniques',\n",
              "  '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.',\n",
              "  'url': 'https://docs.cohere.com/docs/advanced-prompt-engineering-techniques'}]"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "vectorstore.retrieve(\"Prompting by giving examples\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bae81baa",
      "metadata": {},
      "source": [
        "# Run chatbot"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c88017f4",
      "metadata": {},
      "source": [
        "![RAG components - Chatbot](../images/llmu/rag/rag-components-chatbot.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e69fbca9",
      "metadata": {
        "id": "e69fbca9"
      },
      "source": [
        "We can now run the chatbot. For this, we create a generate_chat function which includes the RAG components:\n",
        "- For each user message, we use the endpoint’s search query generation feature to turn the message into one or more queries that are optimized for retrieval. The endpoint can even return no query, which means that a user message can be responded to directly without retrieval. This is done by calling the Chat endpoint with the search_queries_only parameter and setting it as True.\n",
        "- 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.\n",
        "- Finally, all the results from all queries are appended to a list and passed to the Chat endpoint for response generation.\n",
        "- We print the response, together with the citations and the list of document chunks cited, for easy reference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "id": "d2c15a1f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "d2c15a1f",
        "outputId": "8daa9159-338c-45ec-e9ed-830aedcdf0d8"
      },
      "outputs": [],
      "source": [
        "def run_chatbot(message, chat_history=None):\n",
        "    if chat_history is None:\n",
        "        chat_history = []\n",
        "    \n",
        "    # Generate search queries, if any        \n",
        "    response = co.chat(message=message,\n",
        "                        model=\"command-r-plus\",\n",
        "                        search_queries_only=True,\n",
        "                        chat_history=chat_history)\n",
        "    \n",
        "    search_queries = []\n",
        "    for query in response.search_queries:\n",
        "        search_queries.append(query.text)\n",
        "\n",
        "    # If there are search queries, retrieve the documents\n",
        "    if search_queries:\n",
        "        print(\"Retrieving information...\", end=\"\")\n",
        "\n",
        "        # Retrieve document chunks for each query\n",
        "        documents = []\n",
        "        for query in search_queries:\n",
        "            documents.extend(vectorstore.retrieve(query))\n",
        "\n",
        "        # Use document chunks to respond\n",
        "        response = co.chat_stream(\n",
        "            message=message,\n",
        "            model=\"command-r-plus\",\n",
        "            documents=documents,\n",
        "            chat_history=chat_history,\n",
        "        )\n",
        "\n",
        "    else:\n",
        "        response = co.chat_stream(\n",
        "            message=message,\n",
        "            model=\"command-r-plus\",\n",
        "            chat_history=chat_history,\n",
        "        )\n",
        "        \n",
        "    # Print the chatbot response, citations, and documents\n",
        "    chatbot_response = \"\"\n",
        "    print(\"\\nChatbot:\")\n",
        "\n",
        "    for event in response:\n",
        "        if event.event_type == \"text-generation\":\n",
        "            print(event.text, end=\"\")\n",
        "            chatbot_response += event.text\n",
        "        if event.event_type == \"stream-end\":\n",
        "            if event.response.citations:\n",
        "                print(\"\\n\\nCITATIONS:\")\n",
        "                for citation in event.response.citations:\n",
        "                    print(citation)\n",
        "            if event.response.documents:\n",
        "                print(\"\\nCITED DOCUMENTS:\")\n",
        "                for document in event.response.documents:\n",
        "                    print(document)\n",
        "            # Update the chat history for the next turn\n",
        "            chat_history = event.response.chat_history\n",
        "\n",
        "    return chat_history"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dfd03d5f",
      "metadata": {},
      "source": [
        "Here is a sample conversation consisting of a few turns. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "id": "cc071a0b",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Chatbot:\n",
            "Of course! I am here to help. Please go ahead with your question, and I will do my best to assist you."
          ]
        }
      ],
      "source": [
        "# Turn # 1\n",
        "chat_history = run_chatbot(\"Hello, I have a question\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "id": "b84bfdb0",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Retrieving information...\n",
            "Chatbot:\n",
            "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.\n",
            "\n",
            "CITATIONS:\n",
            "start=0 end=19 text='Zero-shot prompting' document_ids=['doc_0']\n",
            "start=29 end=95 text='asking the model to perform a task without providing any examples.' document_ids=['doc_0']\n",
            "start=115 end=133 text='few-shot prompting' document_ids=['doc_0']\n",
            "start=159 end=217 text='model is provided with a few relevant and diverse examples' document_ids=['doc_0']\n",
            "start=246 end=297 text='before asking the specific question to be answered.' document_ids=['doc_0']\n",
            "start=318 end=364 text='steer the model toward a high-quality solution' document_ids=['doc_0']\n",
            "start=369 end=422 text='condition it to the expected response type and style.' document_ids=['doc_0']\n",
            "\n",
            "CITED DOCUMENTS:\n",
            "{'id': 'doc_0', '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.', 'title': 'Advanced Prompt Engineering Techniques', 'url': 'https://docs.cohere.com/docs/advanced-prompt-engineering-techniques'}\n"
          ]
        }
      ],
      "source": [
        "# Turn # 2\n",
        "chat_history = run_chatbot(\"What's the difference between zero-shot and few-shot prompting\", chat_history)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "id": "d60ce722",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Retrieving information...\n",
            "Chatbot:\n",
            "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.\n",
            "\n",
            "CITATIONS:\n",
            "start=23 end=77 text=\"vastly improve the quality of the model's completions.\" document_ids=['doc_2']\n",
            "start=90 end=123 text='few relevant and diverse examples' document_ids=['doc_0']\n",
            "start=130 end=176 text='steer the model toward a high-quality solution' document_ids=['doc_0']\n",
            "start=180 end=236 text='conditioning it to the expected response type and style.' document_ids=['doc_0']\n",
            "\n",
            "CITED DOCUMENTS:\n",
            "{'id': 'doc_2', 'text': 'Advanced Prompt Engineering Techniques\\n\\nSuggest Edits\\n\\nThe previous chapter discussed general rules and heuristics to follow for successfully prompting the Command family of models. Here, we will discuss specific advanced prompt engineering techniques that can in many cases vastly improve the quality of the model’s completions. These include how to give clear and unambiguous instructions, few-shot prompting, chain-of-thought (CoT) techniques, and prompt chaining.', 'title': 'Advanced Prompt Engineering Techniques', 'url': 'https://docs.cohere.com/docs/advanced-prompt-engineering-techniques'}\n",
            "{'id': 'doc_0', '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.', 'title': 'Advanced Prompt Engineering Techniques', 'url': 'https://docs.cohere.com/docs/advanced-prompt-engineering-techniques'}\n"
          ]
        }
      ],
      "source": [
        "# Turn # 3\n",
        "chat_history = run_chatbot(\"How would the latter help?\", chat_history)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "id": "c9faca1e",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Retrieving information...\n",
            "Chatbot:\n",
            "Sorry, I don't have any information about 5G networks. Can I help you with anything else?"
          ]
        }
      ],
      "source": [
        "# Turn # 4\n",
        "chat_history = run_chatbot(\"What do you know about 5G networks?\", chat_history)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d0bf24de",
      "metadata": {},
      "source": [
        "There are a few observations worth pointing out:\n",
        "\n",
        "- Direct response: For user messages that don’t require retrieval (“Hello, I have a question”), the chatbot responds directly without requiring retrieval.\n",
        "- 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.\n",
        "- State management: The endpoint maintains the state of the conversation via the chat_history parameter, for example, by correctly responding to a vague user message such as \"How would the latter help?\"\n",
        "- 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."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "154c95a2",
      "metadata": {},
      "source": [
        "Here are the contents of the chat history."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "id": "3281f3b7",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Chat history:\n",
            "message='Hello, I have a question' tool_calls=None role='USER' \n",
            "\n",
            "message='Of course! I am here to help. Please go ahead with your question, and I will do my best to assist you.' tool_calls=None role='CHATBOT' \n",
            "\n",
            "message=\"What's the difference between zero-shot and few-shot prompting\" tool_calls=None role='USER' \n",
            "\n",
            "message='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.' tool_calls=None role='CHATBOT' \n",
            "\n",
            "message='How would the latter help?' tool_calls=None role='USER' \n",
            "\n",
            "message=\"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.\" tool_calls=None role='CHATBOT' \n",
            "\n",
            "message='What do you know about 5G networks?' tool_calls=None role='USER' \n",
            "\n",
            "message=\"Sorry, I don't have any information about 5G networks. Can I help you with anything else?\" tool_calls=None role='CHATBOT' \n",
            "\n",
            "==================================================\n"
          ]
        }
      ],
      "source": [
        "print(\"Chat history:\")\n",
        "for c in chat_history:\n",
        "    print(c, \"\\n\")\n",
        "print(\"=\"*50)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}