{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hello pgvector: Create, store and query OpenAI embeddings in PostgreSQL using pgvector\n",
    "\n",
    "This notebook will teach you:\n",
    "- How to create embeddings from content using the OpenAI API\n",
    "- How to use PostgreSQL as a vector database and store embeddings data in it using pgvector.\n",
    "- How to use embeddings retrieved from a vector database to augment LLM generation. \n",
    "\n",
    "We'll be using the example of creating a chatbot to answer questions about Timescale use cases, referencing content from the Timescale Developer Q+A blog posts. \n",
    "\n",
    "This is a great first step to building something like chatbot that can reference a company knowledge base or developer docs.\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: This notebook uses a PostgreSQL database with pgvector installed that's hosted on Timescale. You can create your own cloud PostgreSQL database in minutes [at this link](https://console.cloud.timescale.com/signup) to follow along. You can also use a local PostgreSQL database if you prefer.\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration\n",
    "- Signup for an OpenAI Developer Account and create an API Key. See [OpenAI's developer platform](https://platform.openai.com/overview).\n",
    "- Install Python\n",
    "- Install and configure a python virtual environment. We recommend [Pyenv](https://github.com/pyenv/pyenv)\n",
    "- Install the requirements for this notebook using the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "import tiktoken\n",
    "import psycopg2\n",
    "import ast\n",
    "import pgvector\n",
    "import math\n",
    "from psycopg2.extras import execute_values\n",
    "from pgvector.psycopg2 import register_vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run export OPENAI_API_KEY=sk-YOUR_OPENAI_API_KEY...\n",
    "# Get openAI api key by reading local .env file\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "_ = load_dotenv(find_dotenv()) \n",
    "openai.api_key  = os.environ['OPENAI_API_KEY'] "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Create Embeddings\n",
    "First, we'll create embeddings using the OpenAI API on some text we want to augment our LLM with.\n",
    "In this example, we'll use content from the Timescale blog about real world use cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load your CSV file into a pandas DataFrame\n",
    "df = pd.read_csv('blog_posts_data.csv')\n",
    "df.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Calculate cost of embedding data\n",
    "It's usually a good idea to calculate how much creating embeddings for your selected content will cost.\n",
    "We use a number of helper functions to calculate a cost estimate before creating the embeddings to help us avoid surprises.\n",
    "\n",
    "For this toy example, since we're using a small dataset, the total cost will be less than $0.01."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper functions to help us create the embeddings\n",
    "\n",
    "# Helper func: calculate number of tokens\n",
    "def num_tokens_from_string(string: str, encoding_name = \"cl100k_base\") -> int:\n",
    "    if not string:\n",
    "        return 0\n",
    "    # Returns the number of tokens in a text string\n",
    "    encoding = tiktoken.get_encoding(encoding_name)\n",
    "    num_tokens = len(encoding.encode(string))\n",
    "    return num_tokens\n",
    "\n",
    "# Helper function: calculate length of essay\n",
    "def get_essay_length(essay):\n",
    "    word_list = essay.split()\n",
    "    num_words = len(word_list)\n",
    "    return num_words\n",
    "\n",
    "# Helper function: calculate cost of embedding num_tokens\n",
    "# Assumes we're using the text-embedding-ada-002 model\n",
    "# See https://openai.com/pricing\n",
    "def get_embedding_cost(num_tokens):\n",
    "    return num_tokens/1000*0.00002\n",
    "\n",
    "# Helper function: calculate total cost of embedding all content in the dataframe\n",
    "def get_total_embeddings_cost():\n",
    "    total_tokens = 0\n",
    "    for i in range(len(df.index)):\n",
    "        text = df['content'][i]\n",
    "        token_len = num_tokens_from_string(text)\n",
    "        total_tokens = total_tokens + token_len\n",
    "    total_cost = get_embedding_cost(total_tokens)\n",
    "    return total_cost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {},
   "outputs": [],
   "source": [
    "# quick check on total token amount for price estimation\n",
    "total_cost = get_total_embeddings_cost()\n",
    "print(\"estimated price to embed this content = $\" + str(total_cost))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Create smaller chunks of content\n",
    "The OpenAI API has a limit to the maximum amount of tokens it create create an embedding for in a single request. To get around this limit we'll break up our text into smaller chunks. In general its a best practice to create embeddings of a certain size in order to get better retrieval. For our purposes, we'll aim for chunks of around 512 tokens each."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: If you prefer to skip this step, you can use use the provided file: blog_data_and_embeddings.csv which contains the data and embeddings that you'll generate in this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "metadata": {},
   "outputs": [],
   "source": [
    "###############################################################################\n",
    "# Create new list with small content chunks to not hit max token limits\n",
    "# Note: the maximum number of tokens for a single request is 8191\n",
    "# https://platform.openai.com/docs/guides/embeddings/embedding-models\n",
    "###############################################################################\n",
    "# list for chunked content and embeddings\n",
    "new_list = []\n",
    "# Split up the text into token sizes of around 512 tokens\n",
    "for i in range(len(df.index)):\n",
    "    text = df['content'][i]\n",
    "    token_len = num_tokens_from_string(text)\n",
    "    if token_len <= 512:\n",
    "        new_list.append([df['title'][i], df['content'][i], df['url'][i], token_len])\n",
    "    else:\n",
    "        # add content to the new list in chunks\n",
    "        start = 0\n",
    "        ideal_token_size = 512\n",
    "        # 1 token ~ 3/4 of a word\n",
    "        ideal_size = int(ideal_token_size // (4/3))\n",
    "        end = ideal_size\n",
    "        #split text by spaces into words\n",
    "        words = text.split()\n",
    "\n",
    "        #remove empty spaces\n",
    "        words = [x for x in words if x != ' ']\n",
    "\n",
    "        total_words = len(words)\n",
    "        \n",
    "        #calculate iterations\n",
    "        chunks = total_words // ideal_size\n",
    "        if total_words % ideal_size != 0:\n",
    "            chunks += 1\n",
    "        \n",
    "        new_content = []\n",
    "        for j in range(chunks):\n",
    "            if end > total_words:\n",
    "                end = total_words\n",
    "            new_content = words[start:end]\n",
    "            new_content_string = ' '.join(new_content)\n",
    "            new_content_token_len = num_tokens_from_string(new_content_string)\n",
    "            if new_content_token_len > 0:\n",
    "                new_list.append([df['title'][i], new_content_string, df['url'][i], new_content_token_len])\n",
    "            start += ideal_size\n",
    "            end += ideal_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {},
   "outputs": [],
   "source": [
    "openai_client = openai.OpenAI()\n",
    "\n",
    "# Helper function: get embeddings for a text\n",
    "def get_embeddings(text):\n",
    "    response = openai_client.embeddings.create(\n",
    "        model=\"text-embedding-3-small\",\n",
    "        input = text.replace(\"\\n\",\" \")\n",
    "    )\n",
    "    return response.data[0].embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create embeddings for each piece of content\n",
    "for i in range(len(new_list)):\n",
    "    text = new_list[i][1]\n",
    "    embedding = get_embeddings(text)\n",
    "    new_list[i].append(embedding)\n",
    "\n",
    "# Create a new dataframe from the list\n",
    "df_new = pd.DataFrame(new_list, columns=['title', 'content', 'url', 'tokens', 'embeddings'])\n",
    "df_new.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the dataframe with embeddings as a CSV file\n",
    "df_new.to_csv('blog_data_and_embeddings.csv', index=False)\n",
    "# It may also be useful to save as a json file, but we won't use this in the tutorial\n",
    "#df_new.to_json('blog_data_and_embeddings.json') "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Store embeddings with pgvector\n",
    "In this section, we'll store our embeddings and associated metadata. \n",
    "\n",
    "We'll use PostgreSQL as a vector database, with the pgvector extension. \n",
    "\n",
    "You can create a cloud PostgreSQL database for free on [Timescale](https://console.cloud.timescale.com/signup) or use a local PostgreSQL database for this step."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Connect to and configure your vector database\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Timescale database connection string\n",
    "# Found under \"Service URL\" of the credential cheat-sheet or \"Connection Info\" in the Timescale console\n",
    "# In terminal, run: export TIMESCALE_CONNECTION_STRING=postgres://<fill in here>\n",
    "\n",
    "connection_string  = os.environ['TIMESCALE_CONNECTION_STRING'] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect to PostgreSQL database in Timescale using connection string\n",
    "conn = psycopg2.connect(connection_string)\n",
    "cur = conn.cursor()\n",
    "\n",
    "#install pgvector \n",
    "cur.execute(\"CREATE EXTENSION IF NOT EXISTS vector;\")\n",
    "conn.commit()\n",
    "\n",
    "#install pgvectorscale \n",
    "cur.execute(\"CREATE EXTENSION IF NOT EXISTS vectorscale CASCADE;\")\n",
    "conn.commit()\n",
    "\n",
    "# Register the vector type with psycopg2\n",
    "register_vector(conn)\n",
    "\n",
    "# Create table to store embeddings and metadata\n",
    "table_create_command = \"\"\"\n",
    "CREATE TABLE IF NOT EXISTS embeddings (\n",
    "            id bigserial primary key, \n",
    "            title text,\n",
    "            url text,\n",
    "            content text,\n",
    "            tokens integer,\n",
    "            embedding vector(1536)\n",
    "            );\n",
    "            \"\"\"\n",
    "\n",
    "cur.execute(table_create_command)\n",
    "cur.close()\n",
    "conn.commit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optional: Uncomment and execute the following code only if you need to read the embeddings and metadata from the provided CSV file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment and execute this cell only if you need to read the blog data and embeddings from the provided CSV file\n",
    "# Otherwise, skip to next cell\n",
    "'''\n",
    "df = pd.read_csv('blog_data_and_embeddings.csv')\n",
    "titles = df['title']\n",
    "urls = df['url']\n",
    "contents = df['content']\n",
    "tokens = df['tokens']\n",
    "embeds = [list(map(float, ast.literal_eval(embed_str))) for embed_str in df['embeddings']]\n",
    "\n",
    "df_new = pd.DataFrame({\n",
    "    'title': titles,\n",
    "    'url': urls,\n",
    "    'content': contents,\n",
    "    'tokens': tokens,\n",
    "    'embeddings': embeds\n",
    "})\n",
    "'''"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Ingest and store vector data into PostgreSQL using pgvector\n",
    "In this section, we'll batch insert our embeddings and metadata into PostgreSQL and also create an index to help speed up search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 201,
   "metadata": {},
   "outputs": [],
   "source": [
    "register_vector(conn)\n",
    "cur = conn.cursor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remind ourselves of the dataframe structure\n",
    "df_new.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Batch insert embeddings using psycopg2's ```execute_values()```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Batch insert embeddings and metadata from dataframe into PostgreSQL database\n",
    "\n",
    "# Prepare the list of tuples to insert\n",
    "data_list = [(row['title'], row['url'], row['content'], int(row['tokens']), np.array(row['embeddings'])) for index, row in df_new.iterrows()]\n",
    "# Use execute_values to perform batch insertion\n",
    "execute_values(cur, \"INSERT INTO embeddings (title, url, content, tokens, embedding) VALUES %s\", data_list)\n",
    "# Commit after we insert all embeddings\n",
    "conn.commit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sanity check by running some simple queries against the embeddings table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "metadata": {},
   "outputs": [],
   "source": [
    "cur.execute(\"SELECT COUNT(*) as cnt FROM embeddings;\")\n",
    "num_records = cur.fetchone()[0]\n",
    "print(\"Number of vector records in table: \", num_records,\"\\n\")\n",
    "# Correct output should be 129"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print the first record in the table, for sanity-checking\n",
    "cur.execute(\"SELECT * FROM embeddings LIMIT 1;\")\n",
    "records = cur.fetchall()\n",
    "print(\"First record in table: \", records)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create index on embedding column for faster cosine similarity comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an index on the data for faster retrieval\n",
    "# this isn't really needed for 129 vectors, but it shows the usage for larger datasets\n",
    "# Note: always create this type of index after you have data already inserted into the DB\n",
    "\n",
    "# for different tuning suggestions check this: https://github.com/timescale/pgvectorscale?tab=readme-ov-file#tuning\n",
    "cur.execute('CREATE INDEX embedding_idx ON embeddings USING diskann (embedding);')\n",
    "conn.commit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Nearest Neighbor Search using pgvector\n",
    "\n",
    "In this final part of the tutorial, we will query our embeddings table. \n",
    "\n",
    "We'll showcase an example of RAG: Retrieval Augmented Generation, where we'll retrieve relevant data from our vector database and give it to the LLM as context to use when it generates a response to a prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 207,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function: get text completion from OpenAI API\n",
    "# Note max tokens is 4097\n",
    "# Note we're using the latest gpt-3.5-turbo-0613 model\n",
    "def get_completion_from_messages(messages, model=\"gpt-4o\", temperature=0, max_tokens=1000):\n",
    "    response = openai_client.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=messages,\n",
    "        temperature=temperature, \n",
    "        max_tokens=max_tokens, \n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 208,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function: Get top 3 most similar documents from the database\n",
    "def get_top3_similar_docs(query_embedding, conn):\n",
    "    embedding_array = np.array(query_embedding)\n",
    "    # Register pgvector extension\n",
    "    register_vector(conn)\n",
    "    cur = conn.cursor()\n",
    "    # Get the top 3 most similar documents using the KNN <=> operator\n",
    "    cur.execute(\"SELECT content FROM embeddings ORDER BY embedding <=> %s LIMIT 3\", (embedding_array,))\n",
    "    top3_docs = cur.fetchall()\n",
    "    return top3_docs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Define a prompt for the LLM\n",
    "Here we'll define the prompt we want the LLM to provide a reponse to.\n",
    "\n",
    "We've picked an example relevant to the blog post data stored in the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 209,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Question about Timescale we want the model to answer\n",
    "input = \"How is Timescale used in IoT?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to process input with retrieval of most similar documents from the database\n",
    "def process_input_with_retrieval(user_input):\n",
    "    delimiter = \"```\"\n",
    "\n",
    "    #Step 1: Get documents related to the user input from database\n",
    "    related_docs = get_top3_similar_docs(get_embeddings(user_input), conn)\n",
    "\n",
    "    # Step 2: Get completion from OpenAI API\n",
    "    # Set system message to help set appropriate tone and context for model\n",
    "    system_message = f\"\"\"\n",
    "    You are a friendly chatbot. \\\n",
    "    You can answer questions about timescaledb, its features and its use cases. \\\n",
    "    You respond in a concise, technically credible tone. \\\n",
    "    \"\"\"\n",
    "\n",
    "    # Prepare messages to pass to model\n",
    "    # We use a delimiter to help the model understand the where the user_input starts and ends\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_message},\n",
    "        {\"role\": \"user\", \"content\": f\"{delimiter}{user_input}{delimiter}\"},\n",
    "        {\"role\": \"assistant\", \"content\": f\"Relevant Timescale case studies information: \\n {related_docs[0][0]} \\n {related_docs[1][0]} {related_docs[2][0]}\"}   \n",
    "    ]\n",
    "\n",
    "    final_response = get_completion_from_messages(messages)\n",
    "    return final_response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 211,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = process_input_with_retrieval(input)\n",
    "print(input)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 212,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can also ask the model questions about specific documents in the database\n",
    "input_2 = \"Tell me about Edeva and Hopara. How do they use Timescale?\"\n",
    "response_2 = process_input_with_retrieval(input_2)\n",
    "print(input_2)\n",
    "print(response_2)"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
