populate.ipynb

{
    "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 OpenAI's API"
        ]
       },
       {
        "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\""
        ]
       },
       {
        "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",
         "import openai\n",
         "from getpass import getpass\n",
         "\n",
         "\n",
         "import boto3\n",
         "import cassio\n",
         "\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",
         "OPENAI_API_KEY = input(\"Enter your OpenAI API Key:\")"
        ]
       },
       {
        "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": "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=openai.Embedding.retrieve(\"text-embedding-ada-002\"),\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"
        ]
       }
      ],
      "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
     }