{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6715bc2b",
   "metadata": {},
   "source": [
    "# Vector Similarity Astra-Bedrock Search QA Quickstart\n",
    "\n",
    "Set up a simple Question-Answering system with LangChain and Amazon Bedrock, using Astra DB as the Vector Database."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "761d9b70",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "Make sure you have a vector-capable Astra database (get one for free at [astra.datastax.com](https://astra.datastax.com)):\n",
    "\n",
    "- You will be asked to provide the **Database ID** for your Astra DB instance (see [here](https://awesome-astra.github.io/docs/pages/astra/faq/#where-should-i-find-a-database-identifier) for details);\n",
    "- Ensure you have an **Access Token** for your database with role _Database Administrator_ (see [here](https://awesome-astra.github.io/docs/pages/astra/create-token/) for details).\n",
    "\n",
    "Likewise, you will need the credentials to your Amazon Web Services identity, with access to **Amazon Bedrock**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18548e31",
   "metadata": {},
   "source": [
    "## Set up your Python environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "042f832e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install --quiet \\\n",
    "  \"cassio>=0.1.3\" \\\n",
    "  \"langchain==0.0.249\" \\\n",
    "  \"boto3==1.28.62\" \\\n",
    "  \"botocore==1.31.62\" \\\n",
    "  \"cohere==4.37\" \\\n",
    "  \"openai==1.3.7\" \\\n",
    "  \"tiktoken==0.5.2\" \\\n",
    "  \"awscli==1.29.62\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c840842",
   "metadata": {},
   "source": [
    "## Import needed libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b243d1b4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import sys\n",
    "from getpass import getpass\n",
    "\n",
    "\n",
    "import boto3\n",
    "import cassio\n",
    "\n",
    "from langchain.embeddings import BedrockEmbeddings\n",
    "from langchain.llms import Bedrock\n",
    "from langchain.vectorstores import Cassandra\n",
    "from langchain.schema import Document\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.document_loaders import TextLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fccce0c2",
   "metadata": {},
   "source": [
    "## Astra DB Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9b0a92b-f9ce-4810-a8ef-5741b2449b18",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ASTRA_DB_ID = input(\"Enter your Astra DB ID ('0123abcd-'):\")\n",
    "ASTRA_DB_APPLICATION_TOKEN = getpass(\"Enter your Astra DB Token ('AstraCS:...'):\")\n",
    "ASTRA_DB_KEYSPACE = input(\"Enter your keyspace name (optional, default keyspace used if not provided):\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a066f378-8fdb-4d4b-a7b1-bf685fbfd413",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cassio.init(\n",
    "    token=ASTRA_DB_APPLICATION_TOKEN,\n",
    "    database_id=ASTRA_DB_ID,\n",
    "    keyspace=ASTRA_DB_KEYSPACE if ASTRA_DB_KEYSPACE else None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33c674f1",
   "metadata": {},
   "source": [
    "## AWS Credentials Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3219cb02-f955-4fb3-85af-3f149868958a",
   "metadata": {},
   "source": [
    "_Note_: in the following cells you will be asked to explicitly provide the credentials to your AWS account. These are set as environment variables for usage by the subsequent `boto3` calls. Please refer to [boto3's documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html) on the possible ways to supply your credentials in a more production-like environment.\n",
    "\n",
    "In particular, if you are running this notebook in **Amazon SageMaker Studio**, please note that it is sufficient to add the Bedrock policy to your SageMaker role, as outlined at [this link](https://github.com/aws-samples/amazon-bedrock-workshop#enable-aws-iam-permissions-for-bedrock), to access the Bedrock services. In that case you can skip the following three setup cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccbb76ea",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Input your AWS Access Key ID\n",
    "os.environ[\"AWS_ACCESS_KEY_ID\"] = getpass(\"Your AWS Access Key ID:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f93c873d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Input your AWS Secret Access Key\n",
    "os.environ[\"AWS_SECRET_ACCESS_KEY\"] = getpass(\"Your AWS Secret Access Key:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "737249f5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Input your AWS Session Token\n",
    "os.environ[\"AWS_SESSION_TOKEN\"] = getpass(\"Your AWS Session Token:\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4388ac1d",
   "metadata": {},
   "source": [
    "## Set up AWS Bedrock objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d65c46f0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bedrock_runtime = boto3.client(\"bedrock-runtime\", \"us-west-2\")\n",
    "bedrock_embeddings = BedrockEmbeddings(model_id=\"amazon.titan-embed-text-v1\",\n",
    "                                       client=bedrock_runtime)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29d9f48b",
   "metadata": {},
   "source": [
    "## Set up the Vector Store\n",
    "\n",
    "This command will create a suitable table in your database if it does not exist yet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d9f48c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "vector_store = Cassandra(\n",
    "    embedding=bedrock_embeddings,\n",
    "    table_name=\"shakespeare_act5\",\n",
    "    session=None,  # <-- meaning: use the global defaults from cassio.init()\n",
    "    keyspace=None,  # <-- meaning: use the global defaults from cassio.init()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6af24199",
   "metadata": {},
   "source": [
    "## Populate the database\n",
    "\n",
    "Add lines for the text of \"Romeo and Astra\", Scene 5, Act 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dab5114",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# retrieve the text of a scene from act 5 of Romeo and Astra. \n",
    "# Juliet's name was changed to Astra to prevent the LLM from \"cheating\" when providing an answer.\n",
    "! mkdir -p \"texts\"\n",
    "! curl \"https://raw.githubusercontent.com/awesome-astra/docs/main/docs/pages/aiml/aws/bedrock_resources/romeo_astra.json\" \\\n",
    "    --output \"texts/romeo_astra.json\"\n",
    "input_lines = json.load(open(\"texts/romeo_astra.json\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cb6b732",
   "metadata": {},
   "source": [
    "Next, you'll populate the database with the lines from the play.\n",
    "This can take a couple of minutes, please be patient.  In total there are 321 lines.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bae5520",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_documents = []\n",
    "\n",
    "for input_line in input_lines:\n",
    "    if (input_line[\"ActSceneLine\"] != \"\"):\n",
    "        (act, scene, line) = input_line[\"ActSceneLine\"].split(\".\")\n",
    "        location = \"Act {}, Scene {}, Line {}\".format(act, scene, line)\n",
    "        metadata = {\"act\": act, \"scene\": scene, \"line\": line}\n",
    "    else:\n",
    "        location = \"\"\n",
    "        metadata = {}\n",
    "    quote_input = \"{} : {} : {}\".format(location, input_line[\"Player\"], input_line[\"PlayerLine\"])\n",
    "    input_document = Document(page_content=quote_input, metadata=metadata)\n",
    "    input_documents.append(input_document)\n",
    "    \n",
    "print(f\"Adding {len(input_documents)} documents ... \", end=\"\")\n",
    "vector_store.add_documents(documents=input_documents, batch_size=50)\n",
    "print(\"Done.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d162c2d1-188e-43f0-b1c3-342b80641060",
   "metadata": {},
   "source": [
    "## Answer questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5332b838-6c1f-40f4-a29e-2b2d0250f408",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "prompt_template_str = \"\"\"Human: Use the following pieces of context to provide a concise answer to the question at the end.\n",
    "If you don't know the answer, just say that you don't know, don't try to make up an answer.\n",
    "\n",
    "<context>\n",
    "{context}\n",
    "</context\n",
    "\n",
    "Question: {question}\n",
    "\n",
    "Assistant:\"\"\"\n",
    "\n",
    "prompt = PromptTemplate.from_template(prompt_template_str)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea67346b-d5ca-433d-858e-5c21397f9de5",
   "metadata": {
    "tags": []
   },
   "source": [
    "We choose to use the following LLM model (see [this page](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html#model-parameters-general) for more info):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee8b2ee0-f9bf-4ada-8fde-7d917d89c6fa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_id = \"anthropic.claude-v2\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4081fd78-1710-4a42-a942-a14f652c854d",
   "metadata": {},
   "source": [
    "Here the question-answering function is set up, implementing the RAG pattern:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e82ef49-ffec-4429-bcc2-ea09f50333cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "req_accept = \"application/json\"\n",
    "req_content_type = \"application/json\"\n",
    "\n",
    "# This, created from the vector store, will fetch the top relevant documents given a text query\n",
    "retriever = vector_store.as_retriever(search_kwargs={\"k\": 5})\n",
    "\n",
    "def answer_question(question: str, verbose: bool = False) -> str:\n",
    "    if verbose:\n",
    "        print(f\"\\n[answer_question] Question: {question}\")\n",
    "    # Retrieval of the most relevant stored documents from the vector store:\n",
    "    context_docs = retriever.get_relevant_documents(question)\n",
    "    context = \"\\n\".join(doc.page_content for doc in context_docs)\n",
    "    if verbose:\n",
    "        print(\"\\n[answer_question] Context:\")\n",
    "        print(context)\n",
    "    # Filling the prompt template with the current values\n",
    "    llm_prompt_str = prompt.format(\n",
    "        question=question,\n",
    "        context=context,\n",
    "    )\n",
    "    # Invocation of the Amazon Bedrock LLM for text completion -- ultimately obtaining the answer\n",
    "    llm_body = json.dumps({\"prompt\": llm_prompt_str, \"max_tokens_to_sample\": 500})\n",
    "    llm_response = bedrock_runtime.invoke_model(\n",
    "        body=llm_body,\n",
    "        modelId=model_id,\n",
    "        accept=req_accept,\n",
    "        contentType=req_content_type,\n",
    "    )\n",
    "    llm_response_body = json.loads(llm_response[\"body\"].read())\n",
    "    answer = llm_response_body[\"completion\"].strip()\n",
    "    if verbose:\n",
    "        print(f\"\\n[answer_question] Answer: {answer}\\n\")\n",
    "    return answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "599d8856-921e-4bf4-8979-ff54b13de6d5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "my_answer = answer_question(\"Who dies in the story?\")\n",
    "print(\"=\" * 60)\n",
    "print(my_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db71cc88-61e3-42e5-802b-4ba4eaa795b4",
   "metadata": {},
   "source": [
    "Let's take a look at the RAG process piece-wise:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fbeaf2d-f589-4d04-8447-4b876375f5b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "my_answer = answer_question(\"Who dies in the story?\", verbose=True)\n",
    "print(\"=\" * 60)\n",
    "print(my_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "715ba03a-66d1-47fe-9a8e-6b2713ddd0f9",
   "metadata": {},
   "source": [
    "### Interactive QA session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fe0df9c-b24d-4f35-aee4-8bef2dd1e1a6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "user_question = \"\"\n",
    "while True:\n",
    "    user_question = input(\"Enter a question (empty to quit):\").strip()\n",
    "    if user_question:\n",
    "        print(f\"Answer ==> {answer_question(user_question)}\")\n",
    "    else:\n",
    "        print(\"[User, AI exeunt]\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f31c2d97",
   "metadata": {},
   "source": [
    "## Additional resources\n",
    "\n",
    "To learn more about Amazon Bedrock, visit this page: [Introduction to Amazon Bedrock](https://github.com/aws-samples/amazon-bedrock-samples/tree/main/introduction-to-bedrock)."
   ]
  }
 ],
 "metadata": {
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}