notebooks/llmu/co_aws_ch6_rag_bedrock_sm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/cohere-ai/notebooks/blob/main/notebooks/llmu/co_aws_ch6_rag_bedrock_sm.ipynb\"> <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retrieval-Augmented Generation (RAG) Using Cohere on AWS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Large Language Models (LLMs) have proven effective at performing text generation tasks and maintaining the context of a conversation in a chat setting. However, at times, we can encounter a scenario where an LLM hallucinates and provides factually inaccurate responses to a given question. This is especially true in business settings, where companies have proprietary data that an LLM would not have seen during its training phase.\n",
    "\n",
    "Retrieval-augmented generation (RAG) bridges the gap by allowing an LLM to integrate external data sources and use them in its response generation. This significantly minimizes the hallucination issue, making the model's responses more accurate and reliable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install cohere cohere-aws boto3 hnswlib unstructured -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cohere\n",
    "import boto3\n",
    "import cohere_aws\n",
    "from cohere_aws import Client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we set up the clients for Bedrock (to be used for Chat and Embed) and SageMaker (to be used for Rerank) using the same steps as in the previous chapters. Here we name the clients co_br for Bedrock and co_sm for SageMaker.\n",
    "\n",
    "To use Bedrock, we create a BedrockClient by passing the necessary AWS credentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Bedrock client via the native Cohere SDK\n",
    "# Contact your AWS administrator for the credentials\n",
    "\n",
    "co_br = cohere.BedrockClient(\n",
    "    aws_region=\"YOUR_AWS_REGION\",\n",
    "    aws_access_key=\"YOUR_AWS_ACCESS_KEY_ID\",\n",
    "    aws_secret_key=\"YOUR_AWS_SECRET_ACCESS_KEY\",\n",
    "    aws_session_token=\"YOUR_AWS_SESSION_TOKEN\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Later we'll need to create a SageMaker endpoint that exposes access to a Cohere model (Rerank v3 in our case). For this, we’ll use the cohere_aws SDK which makes it easy to set up the endpoint, together with AWS’s boto3 library.\n",
    "\n",
    "Once the endpoint is created (as we’ll walk through later), we can access it using the cohere SDK. To do this, let’s create a SagemakerClient by passing the necessary AWS credentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create SageMaker client via the native Cohere SDK\n",
    "# Contact your AWS administrator for the credentials\n",
    "\n",
    "co_sm = cohere.SagemakerClient(\n",
    "    aws_region=\"YOUR_AWS_REGION\",\n",
    "    aws_access_key=\"YOUR_AWS_ACCESS_KEY_ID\",\n",
    "    aws_secret_key=\"YOUR_AWS_SECRET_ACCESS_KEY\",\n",
    "    aws_session_token=\"YOUR_AWS_SESSION_TOKEN\",\n",
    ")\n",
    "\n",
    "# For creating an endpoint, you need to use the cohere_aws client: Set environment variables with the AWS credentials\n",
    "os.environ['AWS_ACCESS_KEY_ID'] = \"YOUR_AWS_ACCESS_KEY_ID\"\n",
    "os.environ['AWS_SECRET_ACCESS_KEY'] = \"YOUR_AWS_SECRET_ACCESS_KEY\"\n",
    "os.environ['AWS_SESSION_TOKEN'] = \"YOUR_AWS_SESSION_TOKEN\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create SageMaker Endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to create a Rerank SageMaker endpoint by defining the model package Amazon Resource Names (ARN) for the Rerank model. The ARN is an identifying string for a SageMaker resource, and it varies between the regions where a resource is available.\n",
    "\n",
    "Here, we define the Cohere package for the Rerank model and map the model package against each region, which gives the complete ARN for each region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create SageMaker endpoint via the cohere_aws SDK\n",
    "\n",
    "cohere_package = \"cohere-rerank-english-v3-01-d3687e0d2e3a366bb904275616424807\"\n",
    "model_package_map = {\n",
    "    \"us-east-1\": f\"arn:aws:sagemaker:us-east-1:865070037744:model-package/{cohere_package}\",\n",
    "    \"us-east-2\": f\"arn:aws:sagemaker:us-east-2:057799348421:model-package/{cohere_package}\",\n",
    "    \"us-west-1\": f\"arn:aws:sagemaker:us-west-1:382657785993:model-package/{cohere_package}\",\n",
    "    \"us-west-2\": f\"arn:aws:sagemaker:us-west-2:594846645681:model-package/{cohere_package}\",\n",
    "    \"ca-central-1\": f\"arn:aws:sagemaker:ca-central-1:470592106596:model-package/{cohere_package}\",\n",
    "    \"eu-central-1\": f\"arn:aws:sagemaker:eu-central-1:446921602837:model-package/{cohere_package}\",\n",
    "    \"eu-west-1\": f\"arn:aws:sagemaker:eu-west-1:985815980388:model-package/{cohere_package}\",\n",
    "    \"eu-west-2\": f\"arn:aws:sagemaker:eu-west-2:856760150666:model-package/{cohere_package}\",\n",
    "    \"eu-west-3\": f\"arn:aws:sagemaker:eu-west-3:843114510376:model-package/{cohere_package}\",\n",
    "    \"eu-north-1\": f\"arn:aws:sagemaker:eu-north-1:136758871317:model-package/{cohere_package}\",\n",
    "    \"ap-southeast-1\": f\"arn:aws:sagemaker:ap-southeast-1:192199979996:model-package/{cohere_package}\",\n",
    "    \"ap-southeast-2\": f\"arn:aws:sagemaker:ap-southeast-2:666831318237:model-package/{cohere_package}\",\n",
    "    \"ap-northeast-2\": f\"arn:aws:sagemaker:ap-northeast-2:745090734665:model-package/{cohere_package}\",\n",
    "    \"ap-northeast-1\": f\"arn:aws:sagemaker:ap-northeast-1:977537786026:model-package/{cohere_package}\",\n",
    "    \"ap-south-1\": f\"arn:aws:sagemaker:ap-south-1:077584701553:model-package/{cohere_package}\",\n",
    "    \"sa-east-1\": f\"arn:aws:sagemaker:sa-east-1:270155090741:model-package/{cohere_package}\",\n",
    "}\n",
    "\n",
    "region = boto3.Session().region_name\n",
    "if region not in model_package_map.keys():\n",
    "    raise Exception(\"UNSUPPORTED REGION\")\n",
    "\n",
    "model_package_arn = model_package_map[region]\n",
    "\n",
    "co_aws = Client(region_name=region)\n",
    "\n",
    "co_aws.create_endpoint(arn=model_package_arn, endpoint_name=\"my-rerank-v3\", instance_type=\"ml.g5.xlarge\", n_instances=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quick Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We’ll start with a quick example to understand the key aspects of RAG.\n",
    "\n",
    "With RAG, the first step is to define the documents that an LLM will have access to. Here, we have a short list of simple documents. Typically, there is a retrieval process to retrieve the most relevant documents based on a user query, which we’ll cover in the longer walkthrough next. But at this point, let’s assume that these are the only documents and we’ll pass all of them to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = [\n",
    "    {\n",
    "        \"title\": \"Tall penguins\",\n",
    "        \"text\": \"Emperor penguins are the tallest.\"},\n",
    "    {\n",
    "        \"title\": \"Penguin habitats\",\n",
    "        \"text\": \"Emperor penguins only live in Antarctica.\"},\n",
    "    {\n",
    "        \"title\": \"What are animals?\",\n",
    "        \"text\": \"Animals are different from plants.\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have seen how to use the Chat endpoint in the text generation chapter. To use the RAG feature, we simply need to add one additional parameter, documents, to the endpoint call. These are the documents we have defined earlier, which are now available to the model for it to consider utilizing in its response.\n",
    "\n",
    "Let’s now see how the model responds when given the user message, \"What are the tallest living penguins?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "RESPONSE:\n",
      "\n",
      "The tallest living penguins are the Emperor penguins. These penguins only live in Antarctica.\n",
      "\n",
      "CITATIONS:\n",
      "\n",
      "start=4 end=53 text='tallest living penguins are the Emperor penguins.' document_ids=['doc_0']\n",
      "start=69 end=93 text='only live in Antarctica.' document_ids=['doc_1']\n"
     ]
    }
   ],
   "source": [
    "message = \"What are the tallest living penguins?\"\n",
    "\n",
    "response = co_br.chat(message=message,\n",
    "                   documents=documents,\n",
    "                   model=\"cohere.command-r-plus-v1:0\")\n",
    "\n",
    "print(\"\\nRESPONSE:\\n\")\n",
    "print(response.text)\n",
    "    \n",
    "if response.citations:\n",
    "    print(\"\\nCITATIONS:\\n\")           \n",
    "    for citation in response.citations:\n",
    "        print(citation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the response, the model used the documents to inform its answer to the question. For example, the tallest living penguins are emperor penguins part of its response was cited from doc_0, which is the first document in the list containing the text Emperor penguins are the tallest."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A More Comprehensive Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we’ve covered the basics, let’s look at a more comprehensive example of RAG that includes:\n",
    "\n",
    "- Building a retrieval system that includes turning documents into text embeddings and storing them in an index\n",
    "- Building a query generation system that turns user messages into optimized queries for retrieval\n",
    "- Wrapping a user interaction with an LLM in a chat interface\n",
    "- Building a response generation system that’s able to answer different types of queries, such as those that require and don’t require RAG\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let’s import the necessary libraries for this project. This includes hnswlib for the vector library and unstructured for chunking the documents (more details on these later)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we’ll define the documents we’ll use for RAG. We’ll use a few pages from the Cohere documentation that discuss prompt engineering, each in the Python list raw_documents below. Each entry is identified by its title and URL."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "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",
   "metadata": {},
   "source": [
    "# Create Vectorstore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Vectorstore class handles the ingestion of documents into embeddings (or vectors) and the retrieval of relevant documents given a query.\n",
    "\n",
    "It includes a few methods:\n",
    "\n",
    "- load_and_chunk: Loads the raw documents from the URL and breaks them into smaller chunks. We’ll utilize the partition_html method from the unstructured library to perform the chunking.\n",
    "- embed: Generates embeddings of the chunked documents. We use the Embed endpoint available on Bedrock, which uses the cohere.embed-english-v3 model.\n",
    "- index: Indexes the document chunk embeddings to ensure efficient similarity search during retrieval. For this, we’ll use the hnswlib vector library.\n",
    "- retrieve: Uses semantic search to retrieve relevant document chunks from the index, given a query. It involves two steps: first, dense retrieval from the index via the Embed endpoint, and second, a reranking via the Rerank endpoint to boost the search results further."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vectorstore:\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_br.embed(\n",
    "                                texts=texts,\n",
    "                                model=\"cohere.embed-english-v3\",\n",
    "                                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_br.embed(\n",
    "                        texts=[query],\n",
    "                        model=\"cohere.embed-english-v3\",\n",
    "                        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_sm.rerank(\n",
    "                            query=query,\n",
    "                            documents=docs_to_rerank,\n",
    "                            top_n=self.rerank_top_k,\n",
    "                            rank_fields=rank_fields,\n",
    "                            model=\"my-rerank-v3\")\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",
   "metadata": {},
   "source": [
    "# Process Documents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the Vectorstore component is set up, we can process the documents, which will involve chunking, embedding, and indexing. We do this by creating an instance of the Vectorstore and passing the raw documents we defined earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "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",
   "metadata": {},
   "source": [
    "We can test if the retrieval is working by entering a search query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "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": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vectorstore.retrieve(\"Prompting by giving examples\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Chatbot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_chatbot(message, chat_history=None):\n",
    "    \n",
    "    if chat_history is None:\n",
    "        chat_history = []\n",
    "    \n",
    "    # Generate search queries, if any        \n",
    "    response = co_br.chat(message=message,\n",
    "                            search_queries_only=True,\n",
    "                            model=\"cohere.command-r-plus-v1:0\",\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_br.chat(\n",
    "            message=message,\n",
    "            model=\"cohere.command-r-plus-v1:0\",\n",
    "            documents=documents,\n",
    "            chat_history=chat_history)\n",
    "\n",
    "    else:\n",
    "        response = co_br.chat(\n",
    "            message=message,\n",
    "            model=\"cohere.command-r-plus-v1:0\",\n",
    "            chat_history=chat_history)\n",
    "        \n",
    "    # Print the chatbot response, citations, and documents\n",
    "    \n",
    "    print(\"\\nRESPONSE:\\n\")\n",
    "    print(response.text)\n",
    "        \n",
    "    if response.citations:\n",
    "        print(\"\\nCITATIONS:\\n\")           \n",
    "        for citation in response.citations:\n",
    "            print(citation)\n",
    "        print(\"\\nDOCUMENTS:\\n\")           \n",
    "        for document in response.documents:\n",
    "            print(document)\n",
    "            \n",
    "    chat_history = response.chat_history\n",
    "\n",
    "    return chat_history\n",
    "                "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a sample conversation consisting of a few turns. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "RESPONSE:\n",
      "\n",
      "Of course! I am here to help. Please go ahead and ask your question, and I will do my best to provide a helpful response.\n"
     ]
    }
   ],
   "source": [
    "chat_history = run_chatbot(\"Hello, I have a question\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Retrieving information...\n",
      "RESPONSE:\n",
      "\n",
      "Zero-shot prompting is when no examples of the task are provided to the model. On the other hand, few-shot prompting is a technique where a model is given a few examples of the task being performed before asking the specific question to be answered.\n",
      "\n",
      "CITATIONS:\n",
      "\n",
      "start=0 end=19 text='Zero-shot prompting' document_ids=['doc_0']\n",
      "start=28 end=78 text='no examples of the task are provided to the model.' document_ids=['doc_0']\n",
      "start=98 end=116 text='few-shot prompting' document_ids=['doc_0']\n",
      "start=140 end=197 text='model is given a few examples of the task being performed' document_ids=['doc_0']\n",
      "start=205 end=249 text='asking the specific question to be answered.' document_ids=['doc_0']\n",
      "\n",
      "DOCUMENTS:\n",
      "\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": [
    "chat_history = run_chatbot(\"What's the difference between zero-shot and few-shot prompting\", chat_history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Retrieving information...\n",
      "RESPONSE:\n",
      "\n",
      "Few-shot prompting can vastly improve the quality of the model's completions. By providing a few relevant and diverse examples, the model can be steered toward a high-quality solution. These examples condition the model to the expected response type and style.\n",
      "\n",
      "CITATIONS:\n",
      "\n",
      "start=23 end=77 text=\"vastly improve the quality of the model's completions.\" document_ids=['doc_2']\n",
      "start=97 end=126 text='relevant and diverse examples' document_ids=['doc_0']\n",
      "start=145 end=184 text='steered toward a high-quality solution.' document_ids=['doc_0']\n",
      "start=200 end=260 text='condition the model to the expected response type and style.' document_ids=['doc_0']\n",
      "\n",
      "DOCUMENTS:\n",
      "\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": [
    "chat_history = run_chatbot(\"How would the latter help?\", chat_history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Retrieving information...\n",
      "RESPONSE:\n",
      "\n",
      "Sorry, I don't have any information about 5G networks. Is there anything else you would like to ask?\n"
     ]
    }
   ],
   "source": [
    "chat_history = run_chatbot(\"What do you know about 5G networks?\", chat_history)"
   ]
  },
  {
   "cell_type": "markdown",
   "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",
   "metadata": {},
   "source": [
    "Here are the contents of the chat history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "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 and ask your question, and I will do my best to provide a helpful response.' 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 is when no examples of the task are provided to the model. On the other hand, few-shot prompting is a technique where a model is given a few examples of the task being performed before asking the specific question to be answered.' 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. By providing a few relevant and diverse examples, the model can be steered toward a high-quality solution. These examples condition the model 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. Is there anything else you would like to ask?\" 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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrated how to create a RAG application using Cohere Chat and Embed on Amazon Bedrock and Cohere Rerank on Amazon SageMaker. RAG enhances LLMs by enabling them to integrate external data sources and reduce hallucination, resulting in more accurate and reliable responses.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "base"
  },
  "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": 4
}