VectorDB-PGVector/pgvector-and-huggingface.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Beyond Relational Databases: Taming LLM & Transformer Embeddings with PGVector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Objectives\n",
    "\n",
    "- Set up PostgreSQL with the pgvector extension in a Docker container, and create database\n",
    "- Use langchain to add embeddings to database, created with OpenAI's  `text-embedding-ada-002` embedding model\n",
    "- Query the database from langchain to find the most similar embeddings to a given query\n",
    "- Use langchain to add embeddings to database, created with HuggingFace's `all-mpnet-base-v2` embedding model\n",
    "- Use LlamaIndex to add embeddings to database, created with OpenAI's  `text-embedding-ada-002` embedding model\n",
    "- Use sqlalchemy and psycopg2 to add embeddings to database, created with HuggingFace's `all-mpnet-base-v2` embedding model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Install Pre-requisites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install --upgrade pgvector langchain openai psycopg2-binary tiktoken python-dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Vectorizing Text Chunks with Langchain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"<OPENAI_API_KEY>\"\n",
    "\n",
    "## Loading Environment Variables\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary packages\n",
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.embeddings import OpenAIEmbeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read the state_of_the_union.txt file using Langchain's  TextLoader\n",
    "loader = TextLoader('state_of_the_union.txt', encoding='utf-8')\n",
    "documents = loader.load()\n",
    "\n",
    "print(documents)  # prints the document objects\n",
    "print(len(documents))  # 1 - we've only read one file/document into the loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the langchain  RecursiveCharacterTextSplitter object to split this text into chunks.\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=20)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "\n",
    "print(texts)\n",
    "print(len(texts))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "page_content='Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.  \\n\\nLast year COVID-19 kept us apart. This year we are finally together again. \\n\\nTonight, we meet as Democrats Republicans and Independents. But most importantly as Americans. \\n\\nWith a duty to one another to the American people to the Constitution. \\n\\nAnd with an unwavering resolve that freedom will always triumph over tyranny. \\n\\nSix days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated. \\n\\nHe thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined. \\n\\nHe met the Ukrainian people. \\n\\nFrom President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world.' metadata={'source': 'state_of_the_union.txt'}\n"
     ]
    }
   ],
   "source": [
    "print(texts[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\Vinayak Shanawad\\anaconda3\\envs\\pgvector\\Lib\\site-packages\\langchain_core\\_api\\deprecation.py:117: LangChainDeprecationWarning: The class `langchain_community.embeddings.openai.OpenAIEmbeddings` was deprecated in langchain-community 0.1.0 and will be removed in 0.2.0. An updated version of the class exists in the langchain-openai package and should be used instead. To use it run `pip install -U langchain-openai` and import as `from langchain_openai import OpenAIEmbeddings`.\n",
      "  warn_deprecated(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1536\n"
     ]
    }
   ],
   "source": [
    "#  Convert our chunks to embeddings (vectors)\n",
    "embeddings = OpenAIEmbeddings()\n",
    "\n",
    "vector = embeddings.embed_query('Testing the embedding model')\n",
    "\n",
    "print(len(vector))  # 1536 dimensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_vectors = embeddings.embed_documents([t.page_content for t in texts[:5]])\n",
    "\n",
    "print(len(doc_vectors))  # 5 vectors in the output\n",
    "print(doc_vectors[0])    # this will output the first chunk's 1539-dimensional vector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. PGVector for storing Embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.1 PGVector setup\n",
    "\n",
    "- Pull the PGVector database: `docker pull ankane/pgvector`\n",
    "- Start the container with the following command\n",
    "`docker run --name pgvector-demo -e POSTGRES_PASSWORD=test -p 5432:5432 -d ankane/pgvector`\n",
    "- Verify that this is running with: `docker ps`\n",
    "- You can now install a GUI tool such as pgAdmin to inspect the database that is running in the container, or else use psql on the command-line. When connecting, we can specify the host as localhost, and the password as whatever we used in the above command -test, in our case.\n",
    "- We will now create a database, and then add the pgvector extension to that database, with the following SQL commands:\n",
    "    \n",
    "    ```\n",
    "    CREATE DATABASE vector_db;\n",
    "    CREATE EXTENSION vector;\n",
    "    ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.2 Storing texts and embeddings to PGVector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The PGVector Module will try to create a table with the name of the collection.\n",
    "# So, make sure that the collection name is unique and the user has the permission to create a table.\n",
    "\n",
    "from langchain.vectorstores.pgvector import PGVector\n",
    "\n",
    "CONNECTION_STRING = \"postgresql+psycopg2://postgres:test@localhost:5432/vector_db\"\n",
    "COLLECTION_NAME = 'state_of_union_vectors'\n",
    "\n",
    "db = PGVector.from_documents(\n",
    "    embedding=embeddings,\n",
    "    documents=texts,\n",
    "    collection_name=COLLECTION_NAME,\n",
    "    connection_string=CONNECTION_STRING,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(Document(page_content='Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.  \\n\\nLast year COVID-19 kept us apart. This year we are finally together again. \\n\\nTonight, we meet as Democrats Republicans and Independents. But most importantly as Americans. \\n\\nWith a duty to one another to the American people to the Constitution. \\n\\nAnd with an unwavering resolve that freedom will always triumph over tyranny. \\n\\nSix days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated. \\n\\nHe thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined. \\n\\nHe met the Ukrainian people. \\n\\nFrom President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world.', metadata={'source': 'state_of_the_union.txt'}), 0.15866053104400035)\n",
      "\n",
      "(Document(page_content='And tonight I am announcing that we will join our allies in closing off American air space to all Russian flights – further isolating Russia – and adding an additional squeeze –on their economy. The Ruble has lost 30% of its value. \\n\\nThe Russian stock market has lost 40% of its value and trading remains suspended. Russia’s economy is reeling and Putin alone is to blame. \\n\\nTogether with our allies we are providing support to the Ukrainians in their fight for freedom. Military assistance. Economic assistance. Humanitarian assistance. \\n\\nWe are giving more than $1 Billion in direct assistance to Ukraine. \\n\\nAnd we will continue to aid the Ukrainian people as they defend their country and to help ease their suffering.  \\n\\nLet me be clear, our forces are not engaged and will not engage in conflict with Russian forces in Ukraine.  \\n\\nOur forces are not going to Europe to fight in Ukraine, but to defend our NATO Allies – in the event that Putin decides to keep moving west.', metadata={'source': 'state_of_the_union.txt'}), 0.16240880937468893)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Let's do this similarity check in the vector space, and we can use the following Langchain code to do so:\n",
    "query = \"What did the president say about Russia\"\n",
    "similar = db.similarity_search_with_score(query, k=2)\n",
    "\n",
    "for doc in similar:\n",
    "    print(doc, end=\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's get the 1536-dimensional embedding for the above query, with this code:\n",
    "vector = embeddings.embed_query(query)\n",
    "print(vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.3 Working with vectorstore\n",
    "\n",
    "Above, we created a vectorstore from scratch. However, often times we want to work with an existing vectorstore. In order to do that, we can initialize it directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores.pgvector import PGVector\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "\n",
    "embeddings = OpenAIEmbeddings()\n",
    "CONNECTION_STRING = \"postgresql+psycopg2://postgres:test@localhost:5432/vector_db\"\n",
    "COLLECTION_NAME = 'test'\n",
    "\n",
    "store = PGVector(\n",
    "    collection_name=COLLECTION_NAME,\n",
    "    connection_string=CONNECTION_STRING,\n",
    "    embedding_function=embeddings,\n",
    "    \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['3be98acf-c686-11ee-9da3-a4423b5e9371']"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.docstore.document import Document\n",
    "\n",
    "# Add documents\n",
    "store.add_documents([Document(page_content=\"welcome back\")])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. PGVector and LangChain with 768 dim HuggingFace embeddings\n",
    "\n",
    "We can not allow users to store embeddings into a custom table and this may not be a good choice for the data team who would like to create or use existing tables, by default langchain library will store embeddings in the following tables but it can accept the custom embedding size.\n",
    "\n",
    "1. langchain_pg_collection: Store the collection details\n",
    "2. langchain_pg_embedding: Store the embedding details.\n",
    "\n",
    "Discussion reference: https://github.com/langchain-ai/langchain/discussions/17223"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores.pgvector import PGVector\n",
    "import os\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "\n",
    "# embeddings = OpenAIEmbeddings()\n",
    "embeddings = HuggingFaceEmbeddings(model_name=\"all-mpnet-base-v2\")\n",
    "\n",
    "CONNECTION_STRING = \"postgresql+psycopg2://postgres:test@localhost:5432/vector_db\"\n",
    "COLLECTION_NAME = 'huggingface'\n",
    "\n",
    "os.environ[\"PGVECTOR_VECTOR_SIZE\"] = str(768)\n",
    "vectorstore = PGVector(connection_string=CONNECTION_STRING,\n",
    "    embedding_function=embeddings,\n",
    "    collection_name=COLLECTION_NAME,\n",
    "    #pre_delete_collection=True # for testing purposes\n",
    "    )\n",
    "\n",
    "\n",
    "from langchain.docstore.document import Document\n",
    "\n",
    "# Add documents\n",
    "vectorstore.add_documents([Document(page_content=\"How are you today?\")])\n",
    "\n",
    "vector = embeddings.embed_query('How are you today?')\n",
    "vector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. PGVector and LlamaIndex with 1536 dim OpenAI embeddings\n",
    "\n",
    "We can allow users to store embeddings into a custom table by providing the custom embedding size but it creates a table with `data_<table_name>` prefix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install llama_index\n",
    "# !pip install asyncpg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import SimpleDirectoryReader, StorageContext\n",
    "from llama_index.indices.vector_store import VectorStoreIndex\n",
    "from llama_index.vector_stores import PGVectorStore\n",
    "import textwrap\n",
    "import openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Setup OpenAI\n",
    "import os\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"<OPENAI_API_KEY>\"\n",
    "\n",
    "## Loading Environment Variables\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document ID: b7e03f64-4108-459f-827d-87e879731e01\n"
     ]
    }
   ],
   "source": [
    "# Download Data Manually\n",
    "# !mkdir -p 'data/paul_graham/'\n",
    "# !wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'\n",
    "\n",
    "# Loading documents\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham\").load_data()\n",
    "print(\"Document ID:\", documents[0].doc_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the Database\n",
    "import psycopg2\n",
    "\n",
    "connection_string = \"postgresql://postgres:test@localhost:5432\"\n",
    "db_name = \"vector_db\"\n",
    "conn = psycopg2.connect(connection_string)\n",
    "conn.autocommit = True\n",
    "\n",
    "with conn.cursor() as c:\n",
    "    c.execute(f\"DROP DATABASE IF EXISTS {db_name}\")\n",
    "    c.execute(f\"CREATE DATABASE {db_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document ID: 03603262-b343-496b-a602-f764b4187652\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Parsing nodes:   0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Parsing nodes: 100%|██████████| 1/1 [00:00<00:00,  5.87it/s]\n",
      "Generating embeddings: 100%|██████████| 21/21 [00:08<00:00,  2.56it/s]\n"
     ]
    }
   ],
   "source": [
    "from llama_index import SimpleDirectoryReader, StorageContext\n",
    "from llama_index.indices.vector_store import VectorStoreIndex\n",
    "from llama_index.vector_stores import PGVectorStore\n",
    "\n",
    "# Create the index\n",
    "from sqlalchemy import make_url\n",
    "\n",
    "# Loading documents\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham\").load_data()\n",
    "print(\"Document ID:\", documents[0].doc_id)\n",
    "\n",
    "url = make_url(connection_string)\n",
    "vector_store = PGVectorStore.from_params(\n",
    "    database=db_name,\n",
    "    host=url.host,\n",
    "    password=url.password,\n",
    "    port=url.port,\n",
    "    user=url.username,\n",
    "    table_name=\"paul_graham_essay\",\n",
    "    embed_dim=1536,  # openai embedding dimension\n",
    ")\n",
    "\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context, show_progress=True\n",
    ")\n",
    "\n",
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7. PGVector and HuggingFace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install sentence-transformers==2.2.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install sqlalchemy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 7.1 PGVector and HuggingFace with 768 dim embeddings using sqlalchemy Python API\n",
    "\n",
    "The main uses of SQLAlchemy are avoiding the need to write SQL directly, providing a database agnostic ORM, handling connections and transactions, and giving a high level interface while still allowing low level control when needed. The ORM and abstraction from specific database syntax is the major advantage.\n",
    "\n",
    "Make sure we create table first in Pgadmin4:\n",
    "\n",
    "CREATE TABLE HuggingFace (uuid uuid PRIMARY KEY, embedding vector(768), document VARCHAR, cmetadata json);\n",
    "\n",
    "Let's create a CustomPGVector class by referencing langchain so that we can allow users to store embeddings into custom table by providing the custom embedding size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlalchemy\n",
    "from pgvector.sqlalchemy import Vector\n",
    "from sqlalchemy import create_engine, Column, Integer, LargeBinary\n",
    "from sqlalchemy.ext.declarative import declarative_base\n",
    "from sqlalchemy.orm import sessionmaker\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from sqlalchemy.dialects.postgresql import JSON, UUID\n",
    "import uuid\n",
    "\n",
    "from typing import Any, Dict, Iterable, List, Optional, Tuple, Type\n",
    "import logging\n",
    "\n",
    "class CustomPGVector():\n",
    "    def __init__(\n",
    "        self,\n",
    "        model,\n",
    "        connection_string: str,\n",
    "        table_name: str,\n",
    "        embed_dim: int,\n",
    "        logger: Optional[logging.Logger] = None,\n",
    "    ) -> None:\n",
    "        self.table_name = table_name\n",
    "        self.model = model\n",
    "        self.embed_dim = embed_dim\n",
    "        self.connection_string = connection_string\n",
    "        self.logger = logger or logging.getLogger(__name__)\n",
    "        self.__post_init__()\n",
    "    \n",
    "        # Define SQLAlchemy Base and EmbeddingStore classes\n",
    "        self.Base = sqlalchemy.orm.declarative_base()\n",
    "\n",
    "        class EmbeddingStore(self.Base):\n",
    "            __tablename__ = self.table_name\n",
    "\n",
    "            uuid = sqlalchemy.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)\n",
    "            embedding = sqlalchemy.Column(Vector(self.embed_dim))\n",
    "            document = sqlalchemy.Column(sqlalchemy.String, nullable=True)\n",
    "            cmetadata = sqlalchemy.Column(JSON, nullable=True)\n",
    "\n",
    "        self.EmbeddingStore = EmbeddingStore\n",
    "\n",
    "    def __post_init__(\n",
    "        self,\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Initialize the store.\n",
    "        \"\"\"\n",
    "        self._conn = self.connect()\n",
    "\n",
    "\n",
    "    def connect(self) -> sqlalchemy.engine.Connection:\n",
    "        engine = sqlalchemy.create_engine(self.connection_string)\n",
    "        conn = engine.connect()\n",
    "        return conn\n",
    "\n",
    "\n",
    "    def add_texts(\n",
    "        self,\n",
    "        texts: Iterable[str],\n",
    "        metadatas: Optional[List[dict]] = None,\n",
    "        # ids: Optional[List[str]] = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> List[str]:\n",
    "        \"\"\"Run more texts through the embeddings and add to the vectorstore.\n",
    "\n",
    "        Args:\n",
    "            texts: Iterable of strings to add to the vectorstore.\n",
    "            metadatas: Optional list of metadatas associated with the texts.\n",
    "            kwargs: vectorstore specific parameters\n",
    "\n",
    "        Returns:\n",
    "            List of ids from adding the texts into the vectorstore.\n",
    "        \"\"\"    \n",
    "        if not metadatas:\n",
    "            metadatas = [{} for _ in texts]\n",
    "\n",
    "        embeddings = self.embed_documents(list(texts))\n",
    "\n",
    "        Session = sessionmaker(bind=self._conn)\n",
    "        with Session() as session:\n",
    "            uuids = []\n",
    "            for text, metadata, embedding in zip(texts, metadatas, embeddings):\n",
    "                embedding_store = self.EmbeddingStore(\n",
    "                    embedding=embedding,\n",
    "                    document=text,\n",
    "                    cmetadata=metadata\n",
    "                )\n",
    "                session.add(embedding_store)\n",
    "                session.commit()\n",
    "                uuids.append(str(embedding_store.uuid))\n",
    "        return uuids\n",
    "\n",
    "\n",
    "    def embed_documents(self, texts: Iterable[str]):\n",
    "        embeddings = self.model.encode(texts)\n",
    "        return embeddings.tolist()\n",
    "    \n",
    "\n",
    "    def retrieve_embeddings(self, doc_ids: List[str]):\n",
    "        try:\n",
    "            placeholder = ', '.join([\"'\" + e + \"'\" for e in doc_ids])\n",
    "\n",
    "            Session = sessionmaker(bind=self._conn)\n",
    "            with Session() as session:\n",
    "                statement = sqlalchemy.text(f'SELECT embedding FROM {self.table_name} WHERE uuid IN ({placeholder});')\n",
    "                result = session.execute(statement)\n",
    "                embeddings = result.fetchall()\n",
    "                session.commit()\n",
    "        except Exception as e:\n",
    "            self.logger.exception(e)\n",
    "            embeddings = []\n",
    "\n",
    "        return embeddings\n",
    "        \n",
    "\n",
    "    def retrieve_top_k_relevant_docs(self, k: int):\n",
    "        try:\n",
    "            Session = sessionmaker(bind=self._conn)\n",
    "            with Session() as session:\n",
    "                statement = sqlalchemy.text(f'SELECT uuid, document FROM {self.table_name} ORDER BY embedding <-> embedding LIMIT {k};')\n",
    "                result = session.execute(statement)\n",
    "                docs = result.fetchall()\n",
    "                session.commit()\n",
    "        except Exception as e:\n",
    "            self.logger.exception(e)\n",
    "            docs = []\n",
    "\n",
    "        return docs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Store text embeddings into a custom table in PGVector and provide the required embedding size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage\n",
    "model = SentenceTransformer('all-mpnet-base-v2')\n",
    "CONNECTION_STRING = 'postgresql+psycopg2://postgres:test@localhost:5432/vector_db'\n",
    "TABLE_NAME = \"huggingface\"\n",
    "\n",
    "store = CustomPGVector(\n",
    "    model=model,\n",
    "    connection_string=CONNECTION_STRING,\n",
    "    table_name=TABLE_NAME,\n",
    "    embed_dim=768\n",
    ")\n",
    "\n",
    "texts = [\n",
    "\"Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.\",\n",
    "\"Last year COVID-19 kept us apart. This year we are finally together again.\",\n",
    "\"Tonight, we meet as Democrats Republicans and Independents. But most importantly as Americans.\",\n",
    "\"With a duty to one another to the American people to the Constitution.\",\n",
    "\"And with an unwavering resolve that freedom will always triumph over tyranny.\",\n",
    "\"Six days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated.\",\n",
    "\"He thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined.\",\n",
    "\"He met the Ukrainian people.\",\n",
    "\"From President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world.\",\n",
    "\"Groups of citizens blocking tanks with their bodies. Everyone from students to retirees teachers turned soldiers defending their homeland.\",\n",
    "\"In this struggle as President Zelenskyy said in his speech to the European Parliament “Light will win over darkness.” The Ukrainian Ambassador to the United States is here tonight.\",\n",
    "\"Let each of us here tonight in this Chamber send an unmistakable signal to Ukraine and to the world.\",\n",
    "\"Please rise if you are able and show that, Yes, we the United States of America stand with the Ukrainian people.\",\n",
    "\"Throughout our history we’ve learned this lesson when dictators do not pay a price for their aggression they cause more chaos.\",\n",
    "\"They keep moving.\"]\n",
    "\n",
    "uuids = store.add_texts(texts)\n",
    "print(uuids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retrieve embeddings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_id_list = ['1e62faab-5211-4b56-b156-9136a7a00195', '3a6d3b79-368b-4057-9ae2-db4f331a602f']\n",
    "store.retrieve_embeddings(doc_id_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retrieve the top K highest-scoring documents from PGVector:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "docs [(UUID('3a6d3b79-368b-4057-9ae2-db4f331a602f'), 'Last year COVID-19 kept us apart. This year we are finally together again.'), (UUID('3ce34ccf-d78d-4469-99e1-29d1afd00699'), 'Tonight, we meet as Democrats Republicans and Independents. But most importantly as Americans.'), (UUID('29770383-3e7b-45f0-8121-fecb1e2a7a5e'), 'With a duty to one another to the American people to the Constitution.'), (UUID('95326af9-3f25-44d0-9123-fb368698a6a7'), 'And with an unwavering resolve that freedom will always triumph over tyranny.'), (UUID('1e62faab-5211-4b56-b156-9136a7a00195'), 'Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.')]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[(UUID('3a6d3b79-368b-4057-9ae2-db4f331a602f'), 'Last year COVID-19 kept us apart. This year we are finally together again.'),\n",
       " (UUID('3ce34ccf-d78d-4469-99e1-29d1afd00699'), 'Tonight, we meet as Democrats Republicans and Independents. But most importantly as Americans.'),\n",
       " (UUID('29770383-3e7b-45f0-8121-fecb1e2a7a5e'), 'With a duty to one another to the American people to the Constitution.'),\n",
       " (UUID('95326af9-3f25-44d0-9123-fb368698a6a7'), 'And with an unwavering resolve that freedom will always triumph over tyranny.'),\n",
       " (UUID('1e62faab-5211-4b56-b156-9136a7a00195'), 'Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.')]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "store.retrieve_top_k_relevant_docs(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 7.2 PGVector and HuggingFace with 768 dim embeddings using psycopg2 Python API\n",
    "\n",
    "psycopg2 is the new generation psycopg, redesigned to be faster, lighter, and better suited for concurrency. It is the recommended PostgreSQL adapter for current Python versions.\n",
    "\n",
    "Make sure we create table first in Pgadmin4:\n",
    "\n",
    "CREATE TABLE documents (id uuid PRIMARY KEY, document text, embedding vector(768));\n",
    "\n",
    "Let's create a CustomPGVector class so that we can allow users to store and query embeddings into a custom table by providing the custom embedding size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install psycopg2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import psycopg2\n",
    "import uuid\n",
    "\n",
    "from typing import Any, Iterable, List, Optional\n",
    "import logging\n",
    "\n",
    "class CustomPGVector():\n",
    "    def __init__(\n",
    "        self,\n",
    "        model,\n",
    "        connection_string: str,\n",
    "        table_name: str,\n",
    "        embed_dim: int,\n",
    "        logger: Optional[logging.Logger] = None,\n",
    "    ) -> None:\n",
    "        self.table_name = table_name\n",
    "        self.model = model\n",
    "        self.embed_dim = embed_dim\n",
    "        self.connection_string = connection_string\n",
    "        self.logger = logger or logging.getLogger(__name__)\n",
    "        self.__post_init__()\n",
    "\n",
    "    def __post_init__(\n",
    "        self,\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Initialize the store.\n",
    "        \"\"\"\n",
    "        self.cur = self.connect()\n",
    "\n",
    "\n",
    "    def connect(self):\n",
    "        conn = psycopg2.connect(self.connection_string)\n",
    "        conn.autocommit = True\n",
    "        cur = conn.cursor()\n",
    "        return cur\n",
    "\n",
    "\n",
    "    def add_texts(\n",
    "        self,\n",
    "        texts: Iterable[str],\n",
    "        metadatas: Optional[List[dict]] = None,\n",
    "        # ids: Optional[List[str]] = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> List[str]:\n",
    "        \"\"\"Run more texts through the embeddings and add to the vectorstore.\n",
    "\n",
    "        Args:\n",
    "            texts: Iterable of strings to add to the vectorstore.\n",
    "            kwargs: vectorstore specific parameters\n",
    "\n",
    "        Returns:\n",
    "            List of ids from adding the texts into the vectorstore.\n",
    "        \"\"\"    \n",
    "        embeddings = self.embed_documents(list(texts))\n",
    "\n",
    "        uuids = []\n",
    "        for document, embedding in zip(texts, embeddings):\n",
    "            doc_id = str(uuid.uuid4())\n",
    "            self.cur.execute(f'INSERT INTO {self.table_name} (id, tweet, embedding) VALUES (%s, %s, %s) RETURNING id', (doc_id, document, embedding))\n",
    "            uuids.append(self.cur.fetchone()[0])\n",
    "\n",
    "        return uuids\n",
    "\n",
    "\n",
    "    def embed_documents(self, texts: Iterable[str]):\n",
    "        \"\"\"Generate embeddings for text documents.\n",
    "\n",
    "        Args:\n",
    "            texts: Iterable of strings to add to the vectorstore.\n",
    "\n",
    "        Returns:\n",
    "            List of embeddings adding into the vectorstore.\n",
    "        \"\"\"    \n",
    "        embeddings = self.model.encode(texts)\n",
    "        return embeddings.tolist()\n",
    "    \n",
    "\n",
    "    def retrieve_embeddings(self, doc_ids: List[str]):\n",
    "        \"\"\"Retrieve embeddings from vectorstore.\n",
    "\n",
    "        Args:\n",
    "            doc_ids: Iterable of strings to fetch embeddings from the vectorstore.\n",
    "\n",
    "        Returns:\n",
    "            List of embeddings added into the vectorstore.\n",
    "        \"\"\" \n",
    "        try:\n",
    "            placeholder = ', '.join([\"'\" + e + \"'\" for e in doc_ids])\n",
    "            self.cur.execute(f'SELECT embedding FROM {self.table_name} WHERE ID IN ({placeholder});')\n",
    "            embeddings = self.cur.fetchall()\n",
    "        except Exception as e:\n",
    "            self.logger.exception(e)\n",
    "            embeddings = []\n",
    "\n",
    "        return embeddings\n",
    "    \n",
    "\n",
    "    def retrieve_top_k_relevant_docs(self, k: int):\n",
    "        \"\"\"Retrieve top k relevant documents.\n",
    "        We order by embedding <-> embedding which will order by nearest neighbor score.\n",
    "        The <-> operator computes the L2 distance between two vectors. By comparing to itself, it will order by smallest distance (most similar) to largest.\n",
    "\n",
    "        Args:\n",
    "            k: Number of relevant documents to return.\n",
    "\n",
    "        Returns:\n",
    "            List of top k relevant documents from the vectorstore.\n",
    "        \"\"\" \n",
    "        try:\n",
    "            self.cur.execute(f'SELECT id, document FROM {self.table_name} ORDER BY embedding <-> embedding LIMIT {k};')\n",
    "            docs = self.cur.fetchall()\n",
    "        except Exception as e:\n",
    "            self.logger.exception(e)\n",
    "            docs = []\n",
    "\n",
    "        return docs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Store text embeddings into a custom table in PGVector and provide the required embedding size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['01d89557-b546-41f2-a3ed-6dc633a7e383', '41929605-157e-4c88-bf87-f4274261b223', '214fc1ad-cb12-41a8-8c9b-eb4c09db2692', '3ef168ea-1788-4b88-b5f3-aebb5771530f', '75f5fcd4-66bd-418c-9873-403365fefb8b', '4e4c9008-2bd6-445e-abb3-98405c8330b5', 'b3ad6f11-fa80-4cbe-9dcc-cebad2ca6881', '4a4b128a-afed-4433-bf3c-0678336a9245', '6b20bf27-27a9-4123-b5f6-7afdf8af555c', '5fb64d84-72ff-437c-b90d-ac5a8161e207', '4d12a0a7-7d9d-4702-a12b-66a5bf27eb50']\n"
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "# Example usage\n",
    "model = SentenceTransformer('all-mpnet-base-v2')\n",
    "CONNECTION_STRING = 'postgresql://postgres:test@localhost:5432/vector_db'\n",
    "TABLE_NAME = \"DOCUMENTS\"\n",
    "\n",
    "store = CustomPGVector(\n",
    "    model=model,\n",
    "    connection_string=CONNECTION_STRING,\n",
    "    table_name=TABLE_NAME,\n",
    "    embed_dim=768\n",
    "    # table_columns=TABLE_COLUMNS\n",
    ")\n",
    "\n",
    "# texts = [\n",
    "# \"Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.\",\n",
    "# \"Last year COVID-19 kept us apart. This year we are finally together again.\",\n",
    "# \"Tonight, we meet as Democrats Republicans and Independents. But most importantly as Americans.\",\n",
    "# \"With a duty to one another to the American people to the Constitution.\",\n",
    "# \"And with an unwavering resolve that freedom will always triumph over tyranny.\",\n",
    "# \"Six days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated.\",\n",
    "# \"He thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined.\",\n",
    "# \"He met the Ukrainian people.\",\n",
    "# \"From President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world.\",\n",
    "# \"Groups of citizens blocking tanks with their bodies. Everyone from students to retirees teachers turned soldiers defending their homeland.\",\n",
    "# \"In this struggle as President Zelenskyy said in his speech to the European Parliament “Light will win over darkness.” The Ukrainian Ambassador to the United States is here tonight.\",\n",
    "# \"Let each of us here tonight in this Chamber send an unmistakable signal to Ukraine and to the world.\",\n",
    "# \"Please rise if you are able and show that, Yes, we the United States of America stand with the Ukrainian people.\",\n",
    "# \"Throughout our history we’ve learned this lesson when dictators do not pay a price for their aggression they cause more chaos.\",\n",
    "# \"They keep moving.\"]\n",
    "\n",
    "texts = [\n",
    "\"Hello, world!\",\n",
    "\"Nice weather today!\",\n",
    "\"The quick brown fox\",\n",
    "\"I love learning SQL\",\n",
    "\"PostgreSQL is great\",\n",
    "\"Python is my favorite language\",\n",
    "\"Thanks for the explanation\",\n",
    "\"See you later!\",\n",
    "\"I'm learning about NLP\",\n",
    "\"Machine learning is fun\",\"I started liking NLP\"]\n",
    "\n",
    "uuids = store.add_texts(texts)\n",
    "print(uuids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retrieve embeddings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_id_list = ['b4ceeb3a-e9dd-4869-955c-d13f5b9856cc', 'de2ccb67-4815-4045-be98-7138b61095f7']\n",
    "store.retrieve_embeddings(doc_id_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retrieve the top K highest-scoring documents from PGVector:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('de2ccb67-4815-4045-be98-7138b61095f7',\n",
       "  'Last year COVID-19 kept us apart. This year we are finally together again.'),\n",
       " ('18e921d1-0a34-492c-8b36-f839690b9513',\n",
       "  'Tonight, we meet as Democrats Republicans and Independents. But most importantly as Americans.'),\n",
       " ('635e8328-99f9-4ce6-9b0e-b3ff4b8b463d',\n",
       "  'With a duty to one another to the American people to the Constitution.'),\n",
       " ('45ae1d35-0cdd-4a55-86b7-9d56600fe691',\n",
       "  'And with an unwavering resolve that freedom will always triumph over tyranny.'),\n",
       " ('b4ceeb3a-e9dd-4869-955c-d13f5b9856cc',\n",
       "  'Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.')]"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "store.retrieve_top_k_relevant_docs(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 8. References\n",
    "\n",
    "- PGAdmin UI installation - https://www.postgresql.org/ftp/pgadmin/pgadmin4/v8.2/windows/\n",
    "- Postgress pgvector extension - https://www.youtube.com/watch?v=FDBnyJu_Ndg\n",
    "- https://bugbytes.io/posts/vector-databases-pgvector-and-langchain/\n",
    "- Langchain pgvector - https://python.langchain.com/docs/integrations/vectorstores/pgvector\n",
    "- pgvector: https://github.com/pgvector/pgvector\n",
    "- pgvector DockerHub image: https://hub.docker.com/r/ankane/pgvector"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pgvector",
   "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}