Module 8 - Conversational Search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d1d41ec6",
   "metadata": {},
   "source": [
    "# Conversational Search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9afd860b-5997-45a6-869b-06793935c75d",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "\n",
    "In this lab, we leverage Lanchain framework to implement Retrieval Augmented Generation(RAG) solution.RAG retrieves data from outside the language model (non-parametric) and augments the prompts by adding the relevant retrieved data in context. RAG models were introduced by Lewis et al. in 2020 as a model where parametric memory is a pre-trained seq2seq model and the non-parametric memory is a dense vector index of Wikipedia, accessed with a pre-trained neural retriever.\n",
    "\n",
    "In RAG, the external data can come from multiple data sources, such as a document repository, databases, or APIs. The first step is to convert the documents and the user query in the format so they can be compared and relevancy search can be performed. To make the formats comparable for doing relevancy search, a document collection (knowledge library) and the user-submitted query are converted to numerical representation using embedding language models. The embeddings are essentially numerical representations of concept in text. Next, based on the embedding of user query, its relevant text is identified in the document collection by a similarity search in the embedding space. Then the prompt provided by the user is appended with relevant text that was searched and it’s added to the context. The prompt is now sent to the LLM and because the context has relevant external data along with the original prompt, the model output is relevant and accurate.\n",
    "\n",
    "For more informaiton about LangChain RAG, please refere: https://python.langchain.com/docs/use_cases/question_answering/\n",
    "\n",
    "---\n",
    "\n",
    "The lab includes the following steps:\n",
    "1. [Step 1: Initialize](#Step-1:-Initialize)\n",
    "2. [Step 2: Verify deployed endpoint for embedding and content generation model](#Step-2:-Verify-deployed-endpoint-for-embedding-and-content-generation-model)\n",
    "3. [Step 3: Test LLM without context information](#Step-3:-Test-LLM-without-context-information)\n",
    "4. [Step 4: Load documents with LangChain document loader and store vector into OpenSearch](#Step-4:-Load-documents-with-LangChain-document-loader-and-store-vector-into-OpenSearch)\n",
    "5. [Step 5: Retrieval Augmented Generation](#Step-5:-Retrieval-Augmented-Generation)\n",
    "6. [Step 6: Conversational search by memorizing the history](#Step-6:-Conversational-search-by-memorizing-the-history)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03874004-bd68-4fac-a05a-3834986a8e89",
   "metadata": {},
   "source": [
    "## Step 1: Initialize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34490d00-970a-49ac-97f5-1999ff4ca03f",
   "metadata": {},
   "source": [
    "Install required library such as OpenSearch client library, LangChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06184986",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade sagemaker \n",
    "!pip install opensearch-py\n",
    "!pip install unstructured \n",
    "!pip install transformers\n",
    "!pip install langchain==0.3.1\n",
    "!pip install langchain-aws==0.2.1\n",
    "!pip install langchain-community==0.3.1\n",
    "!pip install beautifulsoup4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e422af0-3961-4e99-9f90-f46d3cf78f0b",
   "metadata": {},
   "source": [
    "Initialize SageMaker, Boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "884d0256",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import sagemaker, boto3, json\n",
    "from sagemaker.session import Session\n",
    "\n",
    "sagemaker_session = Session()\n",
    "aws_role = sagemaker_session.get_caller_identity_arn()\n",
    "aws_region = boto3.Session().region_name\n",
    "sess = sagemaker.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad970f25-32cf-4bd8-bfa4-f13e361a5b20",
   "metadata": {},
   "source": [
    "Get Cloud Formation stack output variables\n",
    "\n",
    "We also need to grab some key values from the infrastructure we provisioned using CloudFormation. To do this, we will list the outputs from the stack and store this in \"outputs\" to be used later.\n",
    "\n",
    "You can ignore any \"PythonDeprecationWarning\" warnings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3635d4ec",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "region = aws_region\n",
    "\n",
    "cfn = boto3.client('cloudformation')\n",
    "\n",
    "def get_cfn_outputs(stackname):\n",
    "    outputs = {}\n",
    "    for output in cfn.describe_stacks(StackName=stackname)['Stacks'][0]['Outputs']:\n",
    "        outputs[output['OutputKey']] = output['OutputValue']\n",
    "    return outputs\n",
    "\n",
    "## Setup variables to use for the rest of the demo\n",
    "cloudformation_stack_name = \"semantic-search\"\n",
    "\n",
    "outputs = get_cfn_outputs(cloudformation_stack_name)\n",
    "aos_host = outputs['OpenSearchDomainEndpoint']\n",
    "\n",
    "outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d926ca92-3373-47a7-973e-3136ebfa2370",
   "metadata": {},
   "source": [
    "## Step 2: Verify deployed endpoint for embedding and content generation model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b2378ff-120e-4255-89f1-ea66949df044",
   "metadata": {},
   "source": [
    "Like mentioned in the architecture diagram, we use two LLM in this lab:\n",
    "- Embedding Model: Convert text into vector\n",
    "- Text Generation LLM: Generate content \n",
    "\n",
    "---\n",
    "\n",
    "We will verify the two endpoins are ready before running the lab."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f49640d",
   "metadata": {},
   "source": [
    "### Get endpoint for embedding\n",
    "\n",
    "---\n",
    "This is SageMaker Endpoint with GPT-J 6B parameters model to convert text into vector.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5962e3e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "embedding_endpoint_name=outputs['EmbeddingEndpointName']\n",
    "print(embedding_endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5980c841-860a-46af-a3e2-2bf612c0134d",
   "metadata": {},
   "source": [
    "Verify embedding endpoint is ready"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef415501",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "sm_client = boto3.client(\"sagemaker\", aws_region)\n",
    "\n",
    "describe_embedding_endpoint_response = sm_client.describe_endpoint(EndpointName=embedding_endpoint_name)\n",
    "\n",
    "while describe_embedding_endpoint_response[\"EndpointStatus\"] == 'Creating':\n",
    "    time.sleep(15)\n",
    "    print('.', end='')\n",
    "    describe_embedding_endpoint_response = sm_client.describe_endpoint(EndpointName=embedding_endpoint_name)\n",
    "print('embedding endpoint created')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96b48027",
   "metadata": {},
   "source": [
    "### Get endpoint for content generation\n",
    "\n",
    "---\n",
    "This is SageMaker Endpoint with Falcon 7B parameters model to generate text.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ecca035",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "llm_endpoint_name=outputs['LLMEndpointName']\n",
    "print(llm_endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffe2214c-3487-45a8-babf-d4bf3b3ff021",
   "metadata": {},
   "source": [
    "Verify embedding endpoint is ready"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d029d7a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sm_client = boto3.client(\"sagemaker\", aws_region)\n",
    "\n",
    "describe_llm_endpoint_response = sm_client.describe_endpoint(EndpointName=llm_endpoint_name)\n",
    "\n",
    "while describe_llm_endpoint_response[\"EndpointStatus\"] == 'Creating':\n",
    "    time.sleep(15)\n",
    "    print('.', end='')\n",
    "    describe_llm_endpoint_response = sm_client.describe_endpoint(EndpointName=llm_endpoint_name)\n",
    "print('LLM endpoint created')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b747507b",
   "metadata": {},
   "source": [
    "### Test embedding endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "977a615f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from langchain.embeddings.sagemaker_endpoint import EmbeddingsContentHandler\n",
    "from langchain.embeddings import SagemakerEndpointEmbeddings\n",
    "\n",
    "\n",
    "class TestContentHandler(EmbeddingsContentHandler):\n",
    "    content_type = \"application/json\"\n",
    "    accepts = \"application/json\"\n",
    "\n",
    "    def transform_input(self, prompt: str, model_kwargs={}) -> bytes:\n",
    "        input_str = json.dumps({\"text_inputs\": prompt, **model_kwargs})\n",
    "        return input_str.encode(\"utf-8\")\n",
    "\n",
    "    def transform_output(self, output: bytes) -> str:\n",
    "        response_json = json.loads(output.read().decode(\"utf-8\"))\n",
    "        embeddings = response_json[\"embedding\"]\n",
    "        if len(embeddings) == 1:\n",
    "            return [embeddings[0]]\n",
    "        return embeddings\n",
    "\n",
    "\n",
    "test_content_handler = TestContentHandler()\n",
    "\n",
    "test_embeddings = SagemakerEndpointEmbeddings(\n",
    "    endpoint_name=embedding_endpoint_name,\n",
    "    region_name=aws_region,\n",
    "    content_handler=test_content_handler,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35fb84f3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(test_embeddings.embed_documents([\"Hello World\"])[0][:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa760ac7",
   "metadata": {},
   "source": [
    "### Test LLM endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e86de9ae",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def query_endpoint_with_json_payload(encoded_json, endpoint_name, content_type=\"application/json\"):\n",
    "    client = boto3.client(\"runtime.sagemaker\")\n",
    "    response = client.invoke_endpoint(\n",
    "        EndpointName=endpoint_name, ContentType=content_type, Body=encoded_json\n",
    "    )\n",
    "    return response\n",
    "\n",
    "#method used to parse the inference model's response. we pass it as part of the model's config\n",
    "def parse_response_model(query_response):\n",
    "    model_predictions = json.loads(query_response[\"Body\"].read())\n",
    "    return [gen[\"generated_text\"] for gen in model_predictions]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61f0c128",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "question = \"How to determine shard and data node counts for OpenSearch?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9940794c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "payload = {\n",
    "    \"inputs\": question,\n",
    "    \"parameters\":{\n",
    "        \"max_new_tokens\": 1024,\n",
    "        \"num_return_sequences\": 1,\n",
    "        \"top_k\": 100,\n",
    "        \"top_p\": 0.95,\n",
    "        \"do_sample\": False,\n",
    "        \"return_full_text\": True,\n",
    "        \"temperature\": 0.9\n",
    "    }\n",
    "}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6de27ea",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "query_response = query_endpoint_with_json_payload(\n",
    "    json.dumps(payload).encode(\"utf-8\"), endpoint_name=llm_endpoint_name\n",
    ")\n",
    "\n",
    "generated_texts = parse_response_model(query_response)\n",
    "\n",
    "print(f\"The generated output is: {generated_texts[0]}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "669c4d59-f74f-4636-9e81-754a49bb2fef",
   "metadata": {},
   "source": [
    "## Step 3: Test LLM without context information"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f274d067-e103-4506-850c-1a8fae1d1132",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "To better illustrate why we need retrieval-augmented generation (RAG) based approach to solve the question and anwering problem. LLM can generate answer without context information. However the generated content may be hallucination."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c089f940",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from uuid import uuid4\n",
    "from typing import Dict\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.memory import DynamoDBChatMessageHistory\n",
    "from langchain.memory import ConversationBufferWindowMemory\n",
    "from langchain import PromptTemplate, SagemakerEndpoint\n",
    "from langchain.llms.sagemaker_endpoint import LLMContentHandler\n",
    "from langchain.chains import RetrievalQA\n",
    "\n",
    "\n",
    "class ContentHandler(LLMContentHandler):\n",
    "    content_type = \"application/json\"\n",
    "    accepts = \"application/json\"\n",
    "    \n",
    "    def transform_input(self, prompt: str, model_kwargs={}) -> bytes:\n",
    "        input_str = json.dumps({\"inputs\": prompt, \"parameters\": model_kwargs})\n",
    "        #print(\"Prompt Input:\\n\" + input_str)\n",
    "        return input_str.encode(\"utf-8\")\n",
    "\n",
    "    def transform_output(self, output: bytes) -> str:\n",
    "        response_json = json.loads(output.read().decode(\"utf-8\"))\n",
    "        #print(\"LLM generated text:\\n\" + response_json[0][\"generated_text\"])\n",
    "        return response_json[0][\"generated_text\"]\n",
    "    \n",
    "\n",
    "content_handler = ContentHandler()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4247d17",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "params = {\n",
    "        \"max_length\": 4096,\n",
    "        \"max_new_tokens\": 1024,\n",
    "        \"num_return_sequences\": 1,\n",
    "        \"top_k\": 100,\n",
    "        \"top_p\": 0.95,\n",
    "        \"do_sample\": False,\n",
    "        \"return_full_text\": False,\n",
    "        \"temperature\": 0.9\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb236e6-e242-4d1a-bce3-e2f00409035a",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "llm_hallucination=SagemakerEndpoint(\n",
    "        endpoint_name=llm_endpoint_name,\n",
    "        region_name=aws_region,\n",
    "        model_kwargs=params,\n",
    "        content_handler=content_handler,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccae6fc8-bb14-449e-8eef-0570a6d81093",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Let's directly ask the model a question and see how they respond."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa531974-2eae-4d68-92c0-0a474a80a21d",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Question is:\" + question)\n",
    "generated_result = llm_hallucination(question)\n",
    "print(generated_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71fa6466-9a91-433f-bb70-13daf78df4e2",
   "metadata": {},
   "source": [
    "Ask the same question again to compare the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a16dcaa2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"Question is:\" + question)\n",
    "generated_result = llm_hallucination(question)\n",
    "print(generated_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41c80895",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "The generated result is based on LLM training data, and they are plausible. Moreover, every time the generated content are different. That's the disadvantage of LLM hallucination. \n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cacdeaa9",
   "metadata": {},
   "source": [
    "## Step 4: Load documents with LangChain document loader and store vector into OpenSearch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a128529b-29fc-4e00-9faa-e5de2de70023",
   "metadata": {},
   "source": [
    "Use document loaders to load data from a source as Document's. A Document is a piece of text and associated metadata. For example, there are document loaders for loading a simple .txt file, for loading the text contents of any web page, or even for loading a transcript of a YouTube video.\n",
    "\n",
    "---\n",
    "\n",
    "The following is data flow diagram of loading documents and store vector into OpenSearch.\n",
    "\n",
    "![retriever](./image/module8/document-loader.png)\n",
    "\n",
    "\n",
    "Document loaders expose a \"load\" method for loading data as documents from a configured source. Here, we use `UnstructuredURLLoader` to load OpenSearch best practice web page."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c47a22e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import RecursiveUrlLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size = 1000, chunk_overlap = 100)\n",
    "\n",
    "loader = RecursiveUrlLoader(\n",
    "    \"https://docs.aws.amazon.com/opensearch-service/latest/developerguide/bp.html\",\n",
    "    max_depth=1,\n",
    "    # use_async=False,\n",
    "    # extractor=None,\n",
    "    # metadata_extractor=None,\n",
    "    # exclude_dirs=(),\n",
    "    # timeout=10,\n",
    "    # check_response_status=True,\n",
    "    # continue_on_failure=True,\n",
    "    # prevent_outside=True,\n",
    "    # base_url=None,\n",
    "    # ...\n",
    ")\n",
    "url_texts = loader.load_and_split(text_splitter=text_splitter)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7e4d4df-fbb9-448a-9f43-f33fe7df9dd8",
   "metadata": {},
   "source": [
    "Show one example document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4e53165-93fd-4434-a653-40ba1e98c67d",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "all_splits = url_texts\n",
    "print(url_texts[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57321a3a-db5c-4589-b6a4-6a007154d067",
   "metadata": {},
   "source": [
    "Create an OpenSearch cluster connection.\n",
    "Next, we'll use Python API to set up connection with Amazon Opensearch Service domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ede13e1e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth\n",
    "import boto3\n",
    "import json\n",
    "\n",
    "kms = boto3.client('secretsmanager')\n",
    "aos_credentials = json.loads(kms.get_secret_value(SecretId=outputs['OpenSearchSecret'])['SecretString'])\n",
    "\n",
    "\n",
    "\n",
    "#credentials = boto3.Session().get_credentials()\n",
    "#auth = AWSV4SignerAuth(credentials, region)\n",
    "auth = (aos_credentials['username'], aos_credentials['password'])\n",
    "\n",
    "index_name = 'nlp_pqa'\n",
    "\n",
    "aos_client = OpenSearch(\n",
    "    hosts = [{'host': aos_host, 'port': 443}],\n",
    "    http_auth = auth,\n",
    "    use_ssl = True,\n",
    "    verify_certs = True,\n",
    "    connection_class = RequestsHttpConnection\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95a03874-5a51-456e-b36c-e31771d73ad0",
   "metadata": {
    "tags": []
   },
   "source": [
    "### LangChain embedding endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1671fdfd-0632-49b8-b7f7-05493baa3ece",
   "metadata": {},
   "source": [
    "To build a simiplied QA application with LangChain, we need to wrap up our SageMaker endpoints for embedding model and LLM into `langchain.embeddings.SagemakerEndpointEmbeddings` and `langchain.llms.sagemaker_endpoint.SagemakerEndpoint`. That requires a overwrite methods of `SagemakerEndpointEmbeddings` class to make it compatible with SageMaker embedding mdoel.\n",
    "\n",
    "---\n",
    "\n",
    "Embedding language model is GTP-J, and the endpoint name is `embedding_endpoint_name`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "059dd4b9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from typing import Any, Dict, Iterable, List, Optional, Tuple, Callable\n",
    "import json\n",
    "from langchain.embeddings import SagemakerEndpointEmbeddings\n",
    "from langchain.embeddings.sagemaker_endpoint import EmbeddingsContentHandler\n",
    "from langchain.schema import Document\n",
    "\n",
    "class BulkSagemakerEndpointEmbeddings(SagemakerEndpointEmbeddings):\n",
    "        def embed_documents(\n",
    "            self, texts: List[str], chunk_size: int = 5\n",
    "        ) -> List[List[float]]:\n",
    "            results = []\n",
    "            _chunk_size = len(texts) if chunk_size > len(texts) else chunk_size\n",
    "\n",
    "            for i in range(0, len(texts), _chunk_size):\n",
    "                response = self._embedding_func(texts[i:i + _chunk_size])\n",
    "                results.extend(response)\n",
    "            return results\n",
    "        \n",
    "class EmbeddingContentHandler(EmbeddingsContentHandler):\n",
    "        content_type = \"application/json\"\n",
    "        accepts = \"application/json\"\n",
    "\n",
    "        def transform_input(self, prompt: str, model_kwargs={}) -> bytes:\n",
    "\n",
    "            input_str = json.dumps({\"text_inputs\": prompt, **model_kwargs})\n",
    "            return input_str.encode('utf-8') \n",
    "\n",
    "        def transform_output(self, output: bytes) -> str:\n",
    "\n",
    "            response_json = json.loads(output.read().decode(\"utf-8\"))\n",
    "            embeddings = response_json[\"embedding\"]\n",
    "            if len(embeddings) == 1:\n",
    "                return [embeddings[0]]\n",
    "            return embeddings\n",
    "\n",
    "print(embedding_endpoint_name)\n",
    "embeddings = BulkSagemakerEndpointEmbeddings( \n",
    "            endpoint_name=embedding_endpoint_name,\n",
    "            region_name=aws_region, \n",
    "            content_handler=EmbeddingContentHandler())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57297dfc-3816-4039-8a61-35e0ff1965cc",
   "metadata": {
    "tags": []
   },
   "source": [
    "### OpenSearch vector store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86f0aeb8-d5e5-4f1c-b9c4-432afdbf93b4",
   "metadata": {},
   "source": [
    "Use `OpenSearchVectorSearch` in LangChain to ingest vector into OpenSearch. You can specify more parameters to create kNN index with specified properties. Some parameters like:\n",
    "engine: “nmslib”, “faiss”, “lucene”; default: “nmslib”\n",
    "\n",
    "space_type: “l2”, “l1”, “cosinesimil”, “linf”, “innerproduct”; default: “l2”\n",
    "\n",
    "ef_search: Size of the dynamic list used during k-NN searches. Higher values lead to more accurate but slower searches; default: 512\n",
    "\n",
    "ef_construction: Size of the dynamic list used during k-NN graph creation. Higher values lead to more accurate graph but slower indexing speed; default: 512\n",
    "\n",
    "m: Number of bidirectional links created for each new element. Large impact on memory consumption. Between 2 and 100; default: 16\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec2d425b-0aa0-4e40-9be8-db082b3c5e0c",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.vectorstores import OpenSearchVectorSearch\n",
    "\n",
    "os_domain_ep = 'https://'+aos_host\n",
    "\n",
    "embedding_index_name = 'opensearch_kb_vector'\n",
    "\n",
    "if len(all_splits) > 500:\n",
    "    for i in range(0, len(all_splits), 500):\n",
    "        start = i\n",
    "        end = i+500\n",
    "        if end > len(all_splits):\n",
    "            end = len(all_splits)-1\n",
    "        docs = all_splits[start:end]\n",
    "        OpenSearchVectorSearch.from_documents(\n",
    "            index_name = embedding_index_name,\n",
    "            documents=docs,\n",
    "            embedding=embeddings,\n",
    "            opensearch_url=os_domain_ep,\n",
    "            http_auth=auth\n",
    "        )\n",
    "        print(f\"ingest documents from {start} to {end}\", start, end)\n",
    "else:\n",
    "    OpenSearchVectorSearch.from_documents(\n",
    "            index_name = embedding_index_name,\n",
    "            documents=all_splits,\n",
    "            embedding=embeddings,\n",
    "            opensearch_url=os_domain_ep,\n",
    "            http_auth=auth\n",
    "        )\n",
    "    print(f\"ingest documents\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c369f552-2b2c-49fe-955b-fd814183b720",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "aos_client.indices.get(index=embedding_index_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66bd70c5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#aos_client.indices.delete(index=embedding_index_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93a0d9dd-0350-41ac-a844-2d64ecd11dce",
   "metadata": {},
   "source": [
    "When you use LangChain `OpenSearchVectorSearch` to store embedding with OpenSearch kNN index, you can specify parameters to choose different Approximate Near Neighbour(ANN) algorithms. For more information, please refer OpenSearch kNN documentaion: https://opensearch.org/docs/latest/search-plugins/knn/knn-index/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7610f4d3-1aa7-4a03-a668-90ebd4874d9e",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "customized_embedding_index_name = 'customized_opensearch_kb_vector'\n",
    "\n",
    "OpenSearchVectorSearch.from_documents(\n",
    "            index_name = customized_embedding_index_name,\n",
    "            documents=all_splits,\n",
    "            embedding=embeddings,\n",
    "            opensearch_url=os_domain_ep,\n",
    "            http_auth=auth,\n",
    "            engine=\"faiss\",\n",
    "            space_type=\"innerproduct\",\n",
    "            ef_construction=256,\n",
    "            m=48,\n",
    "        )\n",
    "print(f\"ingest documents into customized knn index\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1aa8d63-6a13-42dd-9123-8d6ef19218a9",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "aos_client.indices.get(index=customized_embedding_index_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c1ecb4d-45ab-4cb7-8fa4-64cc7ef322a7",
   "metadata": {},
   "source": [
    "We can use `OpenSearchVectorSearch` for vector store or we can extend the class to define new fuction to calculate documents relevance score if you want to use relevance score to filter document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8110d127-b60b-45e1-a178-37fd175537ac",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "class SimiliarOpenSearchVectorSearch(OpenSearchVectorSearch):\n",
    "    \n",
    "    def relevance_score(self, distance: float) -> float:\n",
    "        return distance\n",
    "    \n",
    "    def _select_relevance_score_fn(self) -> Callable[[float], float]:\n",
    "        return self.relevance_score\n",
    "    \n",
    "\n",
    "open_search_vector_store = SimiliarOpenSearchVectorSearch(\n",
    "                                    index_name=embedding_index_name,\n",
    "                                    embedding_function=embeddings,\n",
    "                                    opensearch_url=os_domain_ep,\n",
    "                                    http_auth=auth\n",
    "                                    ) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92332ed7-546f-493c-8f42-705275bf70e7",
   "metadata": {},
   "source": [
    "Show the documents which are similiar with question \"How to determine shard and data node counts for OpenSearch?\". Be default, 4 documents are returned. You can specify \"k\" parameter. See the [doc](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.opensearch_vector_search.OpenSearchVectorSearch.html#langchain.vectorstores.opensearch_vector_search.OpenSearchVectorSearch.similarity_search_with_score) for more information.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b644b724",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "docs_ = open_search_vector_store.similarity_search_with_score(question, k=5)\n",
    "\n",
    "print(\"found document number:\" + str(len(docs_)))\n",
    "\n",
    "print(\"opensearch results:\\n\")\n",
    "for doc in docs_:\n",
    "    print(doc)\n",
    "    print(\"\\n-----------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ed0e71e-bf17-4f6e-aaff-be56a874c212",
   "metadata": {},
   "source": [
    "## Step 5: Retrieval Augmented Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db375107-9d31-4e35-acda-a5d20941ca7a",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "To mitigate LLM hallucination, we can provide some context to LLM and let LLM generated answer with the context. The following diagram show RAG data flow:\n",
    "\n",
    "![rag](./image/module8/rag.png)\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f07e9fa",
   "metadata": {},
   "source": [
    "In this lab,  we use OpenSearch vector store as retriever to get similiar documents with query. We can also specify similarity scrore threshhold to return high relevant documents. Use \"k\" to limit how many documents to be returned.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55211dcd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "retriever = open_search_vector_store.as_retriever(\n",
    "    search_type=\"similarity_score_threshold\",\n",
    "    search_kwargs={\n",
    "        'k': 5,\n",
    "        'score_threshold': 0.62\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea2657b1",
   "metadata": {},
   "source": [
    "Define LLM SageMaker endpoint and `RetrievalQA` Chain\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29542ec7-0296-4f3e-b89d-fe9064665960",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "params = {\n",
    "        \"max_length\": 2048,\n",
    "        \"max_new_tokens\": 512,\n",
    "        \"num_return_sequences\": 1,\n",
    "        \"top_k\": 200,\n",
    "        \"top_p\": 0.9,\n",
    "        \"do_sample\": False,\n",
    "        \"return_full_text\": False,\n",
    "        \"temperature\": 0.0001\n",
    "        }\n",
    "\n",
    "llm=SagemakerEndpoint(\n",
    "        endpoint_name=llm_endpoint_name,\n",
    "        region_name=aws_region,\n",
    "        model_kwargs=params,\n",
    "        content_handler=content_handler,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe9a1e14",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "qa = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    retriever=retriever,\n",
    "    chain_type=\"stuff\" #stuff, refine, map_reduce, and map_rerank\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f869cfe8-963f-4053-aa00-22f7ff76a132",
   "metadata": {},
   "source": [
    "Use RAG to generate answer to the same question before. Compare the content generated with RAG and LLM without context.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7feaca4c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"Question is:\" + question)\n",
    "result = qa({\"query\": question})\n",
    "\n",
    "print(\"result:\" + result[\"result\"])\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "774aac71-7220-4205-bf2d-1ed253dd5b1e",
   "metadata": {},
   "source": [
    "### Use customized prompt for RAG"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f479f506-bae2-4573-8126-7b2fa6b42182",
   "metadata": {
    "tags": []
   },
   "source": [
    "---\n",
    "Though the content generated by RAG is much better than the answer without any content. However some of the generated bullets are duplicate. You can improve the generated content by optimizing the prompter to LLM. Following is example of using customized prompt.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0959e688",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "template2 = \"\"\"Answer the question as truthfully as possible by using the provided informaiton in >>CONTEXT<<. If the answer is not contained within the >>CONTEXT<<, respond with \"I can't answer that\".\n",
    "\n",
    ">>CONTEXT<<:\n",
    "{context}\n",
    "\n",
    ">>QUESTION<<:\n",
    "{question}\n",
    "\n",
    ">>Answer<<:\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "prompt_template2 = PromptTemplate(\n",
    "    input_variables=[\"question\", \"context\"],\n",
    "    template=template2\n",
    ")\n",
    "\n",
    "qa_with_prompt = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    retriever=retriever,\n",
    "    chain_type_kwargs={\"prompt\": prompt_template2}\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3de5dec",
   "metadata": {},
   "source": [
    "Run the code qa chain with question. You can set `langchain.debug = True` if you want to see the debug informaiton.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba8517d7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import langchain\n",
    "\n",
    "#langchain.debug = True\n",
    "langchain.debug = False\n",
    "\n",
    "print(\"Question is:\" + question)\n",
    "result = qa_with_prompt({\"query\": question})\n",
    "\n",
    "print(\"\\n### Generated result:\" + result[\"result\"])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5710a4f-5652-4a2d-a126-0ac8e10aceef",
   "metadata": {},
   "source": [
    "### Return source documents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1054077-8a3f-40e1-a5e8-5ef65cbb11c1",
   "metadata": {},
   "source": [
    "You can also return the source documents to help you locate original knowledge base document which are related with the question.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ba727f3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "qa_with_source = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    retriever=retriever,\n",
    "    return_source_documents=True,\n",
    "    chain_type_kwargs={\"prompt\": prompt_template2}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f131847b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"Question is:\" + question)\n",
    "result = qa_with_source({\"query\": question})\n",
    "\n",
    "print(\"result:\" + result[\"result\"])\n",
    "print(\"\\n\\n===========================\")\n",
    "print(\"\\nsource documents:\")\n",
    "for doc in result[\"source_documents\"]:\n",
    "    print(doc)\n",
    "    print(\"---------------------------\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3c88750-e300-45f9-988f-3e238e72f84f",
   "metadata": {},
   "source": [
    "## Step 6: Conversational search by memorizing the history "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60d64dc6-df81-4df2-806a-14cebec9b6ce",
   "metadata": {},
   "source": [
    "### LangChain Memory with Amazon DynamoDB as data store\n",
    "\n",
    "In the above example, you can ask any questions to the system. However there is no relation among the questions. In a typical search system, you may want to implement conversational search. An essential component of a conversation is being able to refer to information introduced earlier in the conversation. LangChain provides a lot of utilities for adding memory to a system. These utilities can be used by themselves or incorporated seamlessly into a chain. In this lab, we use [Amazon DynamoDB](https://aws.amazon.com/dynamodb/) as data store of history message.\n",
    "\n",
    "---\n",
    "The data flow of conversational search with memory is as following:\n",
    "\n",
    "![rag](./image/module8/rag-with-memory.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3824f779-5233-4047-a26a-38e035354445",
   "metadata": {},
   "source": [
    "---\n",
    "Here we create new session and use DynamoDB as backend to store history conversation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25ea68c2",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "ddb_table_name = \"conversation-history-memory\"\n",
    "session_id = str(uuid4())\n",
    "chat_memory = DynamoDBChatMessageHistory(\n",
    "        table_name=ddb_table_name,\n",
    "        session_id=session_id\n",
    "    )\n",
    "\n",
    "messages = chat_memory.messages\n",
    "\n",
    "# Maintains immutable sessions\n",
    "# If previous session was present, create\n",
    "# a new session and copy messages, and \n",
    "# generate a new session_id \n",
    "if messages:\n",
    "    session_id = str(uuid4())\n",
    "    chat_memory = DynamoDBChatMessageHistory(\n",
    "        table_name=ddb_table_name,\n",
    "        session_id=session_id\n",
    "    )\n",
    "    # This is a workaround at the moment. Ideally, this should\n",
    "    # be added to the DynamoDBChatMessageHistory class\n",
    "    try:\n",
    "        messages = messages_to_dict(messages)\n",
    "        chat_memory.table.put_item(\n",
    "            Item={\"SessionId\": session_id, \"History\": messages}\n",
    "        )\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "\n",
    "memory = ConversationBufferMemory(chat_memory=chat_memory, memory_key=\"chat_history\", return_messages=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76dbb51f-2460-4a8e-a482-08be21b02e26",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Create `ConversationalRetrievalChain` to combines the chat history and the question into a standalone question, then looks up relevant documents from the retriever, and finally passes those documents and the question to a question-answering chain to return a response.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfbb0a3e-40a1-4927-8833-5f185410eef8",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.chains import ConversationalRetrievalChain\n",
    "\n",
    "params = {\n",
    "        \"max_length\": 2048,\n",
    "        \"max_new_tokens\": 512,\n",
    "        \"num_return_sequences\": 1,\n",
    "        \"top_k\": 200,\n",
    "        \"top_p\": 0.99,\n",
    "        \"do_sample\": False,\n",
    "        \"return_full_text\": False,\n",
    "        \"temperature\": 0.0001\n",
    "        }\n",
    "\n",
    "llm=SagemakerEndpoint(\n",
    "        endpoint_name=llm_endpoint_name,\n",
    "        region_name=aws_region,\n",
    "        model_kwargs=params,\n",
    "        content_handler=content_handler,\n",
    ")\n",
    "\n",
    "condense_template = \"\"\"system: generate one standalone question.\n",
    "Given the following conversation between <chat-history> and </chat-history> \n",
    "and follow up question between <follow-up-question> and </follow-up-question>, \n",
    "rephrase the follow up question to be a standalone question in its original language. \n",
    "The standalone question will only contains one sentence and it must end with '?'\n",
    "\n",
    "<chat-history>\n",
    "{chat_history}\n",
    "</chat-history>\n",
    "\n",
    "<follow-up-question>\n",
    "{question}\n",
    "</follow-up-question>\n",
    "\n",
    "standalone question:\n",
    "\"\"\"\n",
    "\n",
    "CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(condense_template)\n",
    "\n",
    "\n",
    "qa_with_memory = ConversationalRetrievalChain.from_llm(\n",
    "    llm=llm,\n",
    "    retriever=retriever,\n",
    "    memory=memory,\n",
    "    condense_question_prompt=CONDENSE_QUESTION_PROMPT,\n",
    "    combine_docs_chain_kwargs={\"prompt\": prompt_template2},\n",
    "    verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd000156-1110-4429-b8b9-1a34f1e5ba61",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "For the first question, there is no history. It is just standard RAG process.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0bc1671-79f2-4056-8248-6eff1e3c16fc",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "result = qa_with_memory(question)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee02d8f4-4a23-412f-a76d-41e75037ad9e",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "#print(\"result:\" + str(result))\n",
    "print(\"\\nAnswer:\\n\" + str(result[\"answer\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f5cdd8c-a04d-41fe-99e6-33b531015260",
   "metadata": {},
   "source": [
    "### Second question\n",
    "Try to ask one more question, `ConversationalRetrievalChain` will use the first question, first question's answer and second question as prompt to LLM to generate new question. The prompt to LLM is like following:\n",
    "\n",
    "```python\n",
    "\n",
    "_template = \"\"\"system: generate one standalone question.\n",
    "Given the following conversation between <chat-history> and </chat-history> \n",
    "and follow up question between <follow-up-question> and </follow-up-question>, \n",
    "rephrase the follow up question to be a standalone question in its original language. \n",
    "The standalone question will only contains one sentence and it must end with '?'\n",
    "\n",
    "<chat-history>\n",
    "{chat_history}\n",
    "</chat-history>\n",
    "\n",
    "<follow-up-question>\n",
    "{question}\n",
    "</follow-up-question>\n",
    "\n",
    "standalone question:\n",
    "```\n",
    "\n",
    "After get the new question from LLM, it will search relevant document from OpenSearch vector store and get relevant documents, then combine the new question and relevant documents as prompt to go through RAG process. The prompt to LLM is like following:\n",
    "\n",
    "```python\n",
    "prompt_template = \"\"\"Answer the question as truthfully as possible by using the provided informaiton in >>CONTEXT<<. If the answer is not contained within the >>CONTEXT<<, respond with \"I can't answer that\".\n",
    "\n",
    ">>CONTEXT<<:\n",
    "{context}\n",
    "\n",
    ">>QUESTION<<:\n",
    "{question}\n",
    "\n",
    ">>Answer<<:\n",
    "\"\"\"\n",
    "```\n",
    "\n",
    "In summary, `ConversationalRetrievalChain` will call LLM twice:\n",
    "1. Use history question, history answer and latest question as prompt to generate new question\n",
    "2. Use new question generated in the first step, query relevant documents. Combine relevant documents and new question as prompt to LLM to generate answer.\n",
    "\n",
    "You can also see the verbose message like following:\n",
    "\n",
    "---\n",
    "\n",
    "### First call to LLM:\n",
    "\n",
    "![generate new question](./image/module8/conversation-new-question.png)\n",
    "\n",
    "---\n",
    "\n",
    "### Second call to LLM:\n",
    "\n",
    "![generate final answer](./image/module8/conversation-final-answer.png)\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "762e63fc-6c54-479b-b7c5-8310a16896a3",
   "metadata": {
    "scrolled": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "second_following_question = 'if my data growth is very fast'\n",
    "second_result = qa_with_memory(second_following_question)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8a8099a-a3be-4996-9852-97a6282d85e4",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"second answer:\" + str(second_result[\"answer\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a07d06d-eabf-425a-a9ed-808abf76a271",
   "metadata": {},
   "source": [
    "### Return source document"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a3cc063-ede9-4bed-b780-912a4cbc583c",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "We can also include source documents so that we can know where the content information are from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fb3d864-222a-41f7-a88f-d00754cacda1",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "session_id = str(uuid4())\n",
    "chat_memory = DynamoDBChatMessageHistory(\n",
    "        table_name=ddb_table_name,\n",
    "        session_id=session_id\n",
    "    )\n",
    "\n",
    "memory_for_source = ConversationBufferMemory(chat_memory=chat_memory,memory_key=\"chat_history\")\n",
    "\n",
    "qa_with_memory_and_source = ConversationalRetrievalChain.from_llm(\n",
    "    llm=llm, \n",
    "    retriever=retriever,\n",
    "    condense_question_prompt=CONDENSE_QUESTION_PROMPT,\n",
    "    combine_docs_chain_kwargs={\"prompt\": prompt_template2},\n",
    "    return_source_documents=True, \n",
    "    verbose=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb89fbf6-b500-46ef-af80-006ab3c610e1",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "chat_history = []\n",
    "result = qa_with_memory_and_source({\"question\": question,\"chat_history\": chat_history})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5da474e5-a8d9-4739-a135-5a686146dc7e",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"\\nAnswer:\\n\" + str(result[\"answer\"]))\n",
    "print(\"\\nSource Documents:\\n\")\n",
    "for doc in result[\"source_documents\"]:\n",
    "    print(str(doc))\n",
    "    print(\"--------------------------------\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ef6aa7c-6e9f-41fa-9a7f-fc29b288a1ca",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "chat_history = [(question, result[\"answer\"])]\n",
    "second_query = second_following_question\n",
    "second_result = qa_with_memory_and_source({\"question\": second_query, \"chat_history\": chat_history})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8965962-cd1d-481b-bed2-e3f475252814",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"\\nAnswer:\\n\" + str(second_result[\"answer\"]))\n",
    "print(\"\\nSource Documents:\\n\")\n",
    "for doc in second_result[\"source_documents\"]:\n",
    "    print(str(doc))\n",
    "    print(\"--------------------------------\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "009b53e7-4d13-4251-a580-5f5bff3d9371",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}