notebooks/03_build_pgvector_db.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supercharging Vector Similarity Search with Amazon Aurora and pgvector\n",
    "In this Jupyter Notebook, you'll explore how to store vector embeddings in a vector database using [Amazon Aurora](https://aws.amazon.com/es/rds/aurora/) and the pgvector extension. This approach is particularly useful for applications that require efficient similarity searches on high-dimensional data, such as natural language processing, image recognition, and recommendation systems.\n",
    "\n",
    "[Amazon Aurora](https://aws.amazon.com/es/rds/aurora/) is a fully managed relational database service provided by Amazon Web Services (AWS). It is compatible with PostgreSQL and supports the [pgvector](https://github.com/pgvector/pgvector) extension, which introduces a 'vector' data type and specialized query operators for vector similarity searches. The pgvector extension utilizes the ivfflat indexing mechanism to expedite these searches, allowing you to store and index up to 16,000 dimensions, while optimizing search performance for up to 2,000 dimensions.\n",
    "\n",
    "For developers and data engineers with experience in relational databases and PostgreSQL, Amazon Aurora with pgvector offers a powerful and familiar solution for managing vector datastores, especially when dealing with structured datasets. Alternatively, Amazon Relational Database Service (RDS) for PostgreSQL is also a suitable option, particularly if you require specific PostgreSQL versions.\n",
    "\n",
    "Both Amazon Aurora and Amazon RDS for PostgreSQL offer horizontal scaling capabilities for read queries, with a maximum of 15 replicas. Additionally, Amazon Aurora PostgreSQL provides a Serverless v2 option, which automatically scales compute and memory resources based on your application's demand, simplifying operations and capacity planning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get started with storing embeddings in a vector database using Amazon Aurora and pgvector, follow these steps:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[langchain_postgres](https://github.com/langchain-ai/langchain-postgres?tab=readme-ov-file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install psycopg\n",
    "#!pip install langchain_postgres"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**1- Set up an Amazon Aurora instance:** Ensure that you have an Amazon Aurora instance configured and running. Add all the necessary connection details, such as the endpoint, username, and password, to your application's environment variables or a .env file.\n",
    "\n",
    "> [Follow steps here. ](https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/Aurora.CreateInstance.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order for you to connect to the Aurora instance from your computer using this notebook, you must allow public access.\n",
    "\n",
    "> Learn more in [How do I configure a provisioned Amazon Aurora DB cluster to be publicly accessible?](https://repost.aws/knowledge-center/aurora-mysql-connect-outside-vpc)\n",
    "\n",
    "![Aurora public](aurora_public.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order for use the endpoint to run SQL statements without managing connections..\n",
    "\n",
    "> Learn more in [Using RDS Data API](https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/data-api.html)\n",
    "\n",
    "![RDS Data API](enable_RDS_api.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You are going to need the Database Name to enable the connection.\n",
    "\n",
    "![DataBase_name](database_name.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And add a new rule with your IP in the Inbound [security group](https://docs.aws.amazon.com/vpc/latest/userguide/vpc-security-groups.html)\n",
    "\n",
    "![Security Group](security_group.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Get the host name](https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/Aurora.Connecting.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PGVECTOR_DRIVER='psycopg2'\n",
    "PGVECTOR_USER='<<Username>>'\n",
    "PGVECTOR_PASSWORD='<<Password>>'\n",
    "PGVECTOR_HOST='<<Aurora DB cluster host>>'\n",
    "PGVECTOR_PORT=5432\n",
    "PGVECTOR_DATABASE='<<DBName>>'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "driver=os.getenv(\"PGVECTOR_DRIVER\"),\n",
    "user=os.getenv(\"PGVECTOR_USER\"),\n",
    "password=os.getenv(\"PGVECTOR_PASSWORD\"),\n",
    "host=os.getenv(\"PGVECTOR_HOST\"),\n",
    "port=os.getenv(\"PGVECTOR_PORT\"),\n",
    "database=os.getenv(\"PGVECTOR_DATABASE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import psycopg\n",
    "\n",
    "connection = f\"postgresql+psycopg://{user}:{password}@{host}:{port}/{database}\" \n",
    "# Establish the connection to the database\n",
    "conn = psycopg.connect(\n",
    "    conninfo = f\"postgresql://{user}:{password}@{host}:{port}/{database}\"\n",
    ")\n",
    "# Create a cursor to run queries\n",
    "cur = conn.cursor()\n",
    "# Create a vectorstore\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**2- Enable the [pgvector extension](https://docs.aws.amazon.com/AmazonRDS/latest/AuroraPostgreSQLReleaseNotes/AuroraPostgreSQL.Extensions.html?sc_channel=el&sc_campaign=genai&sc_geo=mult&sc_country=mult&sc_outcome=acq&sc_content=vector-embeddings-and-rag-demystified-2):** Once connected to your Aurora instance, enable the pgvector extension by running the following SQL command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cur.execute(\"CREATE EXTENSION vector;\")\n",
    "conn.commit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**3 -Create a table to store embeddings:** Define a table schema to store your vector embeddings and any associated metadata. \n",
    "\n",
    "This table includes columns for a unique identifier (id), the original text (text), and the vector embedding (embedding) with a dimensionality of 1536."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table_name = \"embeddings\"\n",
    "query = f\"\"\"CREATE TABLE {table_name} (\n",
    "    id SERIAL PRIMARY KEY,\n",
    "    text TEXT,\n",
    "    embedding VECTOR(1536)\n",
    ");\"\"\"\n",
    "cur.execute(query)\n",
    "conn.commit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. **Insert embeddings into the table using Langchain**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_aws import BedrockEmbeddings # to create embeddings for the documents.\n",
    "from langchain_experimental.text_splitter import SemanticChunker # to split documents into smaller chunks.\n",
    "from langchain.docstore.document import Document\n",
    "import boto3\n",
    "from langchain_postgres import PGVector\n",
    "from langchain_postgres.vectorstores import PGVector\n",
    "from langchain_community.document_loaders import PyPDFLoader\n",
    "from langchain_community.chat_models import BedrockChat\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.callbacks import StdOutCallbackHandler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrock_client              = boto3.client(\"bedrock-runtime\") \n",
    "bedrock_embeddings          = BedrockEmbeddings(model_id=\"amazon.titan-embed-text-v1\",client=bedrock_client)\n",
    "bedrock_embeddings_image = BedrockEmbeddings(model_id=\"amazon.titan-embed-image-v1\",client=bedrock_client)\n",
    "llm = BedrockChat(model_id=\"anthropic.claude-3-sonnet-20240229-v1:0\", client=bedrock_client)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to create vector store\n",
    "def create_vectorstore(embeddings,collection_name,conn):\n",
    "                       \n",
    "    vectorstore = PGVector(\n",
    "        embeddings=embeddings,\n",
    "        collection_name=collection_name,\n",
    "        connection=conn,\n",
    "        use_jsonb=True,\n",
    "    )\n",
    "    return vectorstore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Retrieve information using Amazon Bedrock\n",
    "def retrieve_information(llm, vectordb,query):\n",
    "    # Set up the retrieval chain with the language model and database retriever\n",
    "    chain = RetrievalQA.from_chain_type(\n",
    "                                            llm=llm,\n",
    "                                            retriever=vectordb.as_retriever(),\n",
    "                                            verbose=True\n",
    "                                        )\n",
    "\n",
    "    # Initialize the output callback handler\n",
    "    handler = StdOutCallbackHandler()\n",
    "\n",
    "    # Run the retrieval chain with a query\n",
    "    chain_value = chain.run(\n",
    "                query,\n",
    "                callbacks=[handler]\n",
    "            )\n",
    "    return chain_value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PDF File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and process PDF documents\n",
    "file_name = \"Amazon_Bedrock_User_Guide.pdf\"\n",
    "path_file = \"demo-files\"\n",
    "file_path = f\"{path_file}/{file_name}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_and_split_pdf_semantic(file_path, embeddings):\n",
    "    text_splitter = SemanticChunker(embeddings, breakpoint_threshold_amount= 80)\n",
    "    loader = PyPDFLoader(file_path)\n",
    "    docs = loader.load_and_split(text_splitter)\n",
    "    print(f\"docs:{len(docs)}\")\n",
    "    return docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = load_and_split_pdf_semantic(file_path, bedrock_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "collection_name_text = \"text_collection\"\n",
    "vectorstore = create_vectorstore(bedrock_embeddings,collection_name_text,connection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add documents to the vectorstore\n",
    "vectorstore.add_documents(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vector retriever\n",
    "More information: \n",
    "- [pgvector](https://python.langchain.com/docs/integrations/vectorstores/pgvector/)\n",
    "- [Retrievers](https://python.langchain.com/docs/modules/data_connection/retrievers/)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectorstore.similarity_search(\"what is a prompt\", k=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectorstore.similarity_search_with_relevance_scores(\"what is a prompt\", k=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrieve information using Amazon Bedrock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"what is aprompt?\"\n",
    "response = retrieve_information(llm, vectorstore,query) \n",
    "print(response)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Learn more: \n",
    "- [Leverage pgvector and Amazon Aurora PostgreSQL for Natural Language Processing, Chatbots and Sentiment Analysis](https://aws.amazon.com/es/blogs/database/leverage-pgvector-and-amazon-aurora-postgresql-for-natural-language-processing-chatbots-and-sentiment-analysis/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import base64\n",
    "import os\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#calls Bedrock to get a vector from either an image, text, or both\n",
    "def get_multimodal_vector(input_image_base64=None, input_text=None):\n",
    "    request_body = {}\n",
    "    if input_text:\n",
    "        request_body[\"inputText\"] = input_text\n",
    "        \n",
    "    if input_image_base64:\n",
    "        request_body[\"inputImage\"] = input_image_base64\n",
    "    \n",
    "    body = json.dumps(request_body)\n",
    "    response = bedrock_client.invoke_model(\n",
    "    \tbody=body, \n",
    "    \tmodelId=\"amazon.titan-embed-image-v1\", \n",
    "    \taccept=\"application/json\", \n",
    "    \tcontentType=\"application/json\"\n",
    "    )\n",
    "\n",
    "    response_body = json.loads(response.get('body').read())\n",
    "    \n",
    "    embedding = response_body.get(\"embedding\")\n",
    "    \n",
    "    return embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Función para convertir una imagen a base64\n",
    "def image_to_base64(image_path):\n",
    "    with open(image_path, \"rb\") as image_file:\n",
    "        encoded_string = base64.b64encode(image_file.read()).decode(\"utf-8\")\n",
    "    vector = get_multimodal_vector(input_image_base64=encoded_string)\n",
    "    return encoded_string, vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_size_image(file_path):\n",
    "    # Maximum image size supported is 2048 x 2048 pixel\n",
    "    image = Image.open(file_path) #open image\n",
    "    width, height = image.size # Get the width and height of the image in pixels\n",
    "    if width > 2048 or height > 2048:\n",
    "        print(f\"Big File:{file_path} , width: {width}, height {height} px\")\n",
    "        dif_width = width - 2048\n",
    "        dif_height = height - 2048\n",
    "        if dif_width > dif_height:\n",
    "            ave = 1-(dif_width/width)\n",
    "            new_width = int(width*ave)\n",
    "            new_height = int(height*ave)\n",
    "        else:\n",
    "            ave = 1-(dif_height/height)\n",
    "            new_width = int(width*ave)\n",
    "            new_height = int(height*ave)\n",
    "        print(f\"New file: {file_path} , width: {new_width}, height {new_height} px\")\n",
    "\n",
    "        new_image = image.resize((new_width, new_height))\n",
    "        # Save New image\n",
    "        new_image.save(file_path)\n",
    " \n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image_vectors_from_directory(path_name):\n",
    "    documents = []\n",
    "    embeddings = []\n",
    "    for folder in os.walk(path_name):\n",
    "        #print(f'In {folder[0]} are {len(folder[2])} folder:')\n",
    "        for fichero in folder[2]:\n",
    "            if fichero.endswith('.jpg'):\n",
    "                file_path = os.path.join(folder[0], fichero)\n",
    "                #print(file_path)\n",
    "                check_size_image(file_path)\n",
    "                image_base64, image_embedding = image_to_base64(file_path)\n",
    "                documents.append({\"page_content\": image_base64, \"file_path\": file_path})\n",
    "                embeddings.append(image_embedding)\n",
    "            else:\n",
    "                print(\"no a .jpg file: \", fichero)\n",
    "\n",
    "    return documents, embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_file = \"animals/animals\"\n",
    "path_name = f\"{path_file}\"\n",
    "documents, embeddings = get_image_vectors_from_directory(path_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#image_vectorstore.drop_tables()\n",
    "collection_name_image = \"image_collection\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_vectorstore = create_vectorstore(bedrock_embeddings_image,collection_name_image,connection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = [d.get(\"file_path\") for d in documents]\n",
    "metadata = [{\"file_path\": d.get(\"file_path\")} for d in documents]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_vectorstore.add_embeddings(embeddings=embeddings, texts=texts, metadata=metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "similitary = image_vectorstore.similarity_search_with_relevance_scores(\"a woodpecker\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in similitary:\n",
    "    print(\"Path:\",n[0].page_content)\n",
    "    print(\"Relevance Score:\",n[1])\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrieve information using Amazon Bedrock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"a woodpecker\"\n",
    "response = retrieve_information(llm, image_vectorstore,query) \n",
    "print(response)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retriever by Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_path = \"animals/animals/whale/3e8b0a420a.jpg\"\n",
    "image_base64, image_embedding = image_to_base64(image_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "similitary_by_vector = image_vectorstore.similarity_search_with_score_by_vector(image_embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in similitary_by_vector:\n",
    "    print(\"Path:\",n[0].page_content)\n",
    "    print(\"Difference between vectors:\",n[1]) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using RDS Data API\n",
    "https://github.com/aws-samples/amazon-rds-data-api-demo/blob/master/src/main/python/lambda_function_postgres.py\n",
    "\n",
    "https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/rds-data.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "cluster_arn = 'arn:aws:rds:us-east-1:0000000000:cluster:my-cluster' \n",
    "secret_arn = 'arn:aws:secretsmanager:us-east-1:0000000000:secret:my-secret'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def call_rds_data_api(timestamp, message):\n",
    "    rds_data = boto3.client('rds-data')\n",
    "\n",
    "    sql = \"\"\"\n",
    "          INSERT INTO sample_table(received_at, message)\n",
    "          VALUES(TO_TIMESTAMP(:time, 'YYYY-MM-DD HH24:MI:SS'), :message)\n",
    "          \"\"\"\n",
    "\n",
    "    param1 = {'name':'time', 'value':{'stringValue': timestamp}}\n",
    "    param2 = {'name':'message', 'value':{'stringValue': message}}\n",
    "    param_set = [param1, param2]\n",
    " \n",
    "    response = rds_data.execute_statement(\n",
    "        resourceArn = cluster_arn, \n",
    "        secretArn = secret_arn, \n",
    "        database = 'tutorial', \n",
    "        sql = sql,\n",
    "        parameters = param_set)\n",
    "    \n",
    "    print(str(response));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rds_data = boto3.client('rds-data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Delete vectorDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectorstore.drop_tables()\n",
    "image_vectorstore.drop_tables()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}