{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple RAG System using pgai, pgai vectorizer, OpenAI's text-embedding-3-small & GPT-4o mini in PostgreSQL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Storing OpenAI API key & database connection string\n",
    "\n",
    "* **OpenAI API Key**:  navigate to [OpenAI Platform](https://platform.openai.com/api-keys) to get the key. You will need to [sign up for an OpenAI Developer Account](https://auth.openai.com/authorize?issuer=auth0.openai.com&client_id=DRivsnm2Mu42T3KOpqdtwB3NYviHYzwD&audience=https%3A%2F%2Fapi.openai.com%2Fv1&redirect_uri=https%3A%2F%2Fplatform.openai.com%2Fauth%2Fcallback&device_id=79ac50b7-1d7b-4d25-bf75-b47c8b74a76d&screen_hint=signup&max_age=0&scope=openid+profile+email+offline_access&response_type=code&response_mode=query&state=OWsydkIzLTZEcThrUX51dmgwN09qUHUxX09RTG45c2tMZH45OXpqfjBxWA%3D%3D&nonce=YVVkaGV0Szl0anV1MHp5Yk5oUVZtNDhlMmdLTWtCTUFzc1c3Z2NhdVZxNA%3D%3D&code_challenge=tvS6MJpVf11ZFLdXAlRREubma_WKaDkgfqLw1ZfY8jg&code_challenge_method=S256&auth0Client=eyJuYW1lIjoiYXV0aDAtc3BhLWpzIiwidmVyc2lvbiI6IjEuMjEuMCJ9&flow=control), if you don't have one yet. \n",
    "\n",
    "* **PostgreSQL database connection string**: navigate to the [Timescale console](https://console.cloud.timescale.com/signup/?utm_source=blog&utm_medium=website&utm_campaign=vectorlaunch&utm_content=automate-embeddings-cta), create a database service, and get the database connection string.\n",
    "\n",
    "Then create a `.env` file to store them in your python environment as [best practice](https://help.openai.com/en/articles/5112595-best-practices-for-api-key-safety?ref=timescale.com). Add them to this file in the following format:\n",
    "\n",
    "`export OPENAI_API_KEY='your_OPENAI_API_KEY'`\n",
    "\n",
    "`export DATABASE_CONNECTION_STRING='postgres://tsdbadmin:db_password@host:port/tsdb?sslmode=require'`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install Python Libraries\n",
    "\n",
    "Let's install the python libraries that we will use in this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install pandas psycopg2-binary python-dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the dataset provided in the same folder `product-catalog.csv`. We divide into two subsets to demonstrate the following features of pgai vectorizer:\n",
    "\n",
    "- embedding creation using `subset_1`. \n",
    "- updating the embeddings with source data change using the `subset_2`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "DATASET_URL = 'https://raw.githubusercontent.com/timescale/vector-cookbook/refs/heads/main/intro_pgaivectorizer_rag/product_catalog.csv'\n",
    "data = pd.read_csv(DATASET_URL)\n",
    "\n",
    "products = [\n",
    "    {\"name\": row[\"name\"], \"description\": row[\"description\"]}\n",
    "    for _, row in data.iterrows()\n",
    "]\n",
    "\n",
    "subset_1 = products[:250]\n",
    "subset_2 = products[250:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to the database"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first extract the database connection string and OpenAI API key from the environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "OPENAI_API_KEY = os.environ[\"OPENAI_API_KEY\"]\n",
    "DATABASE_CONNECTION_STRING = os.environ[\"DATABASE_CONNECTION_STRING\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now connect to the database service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import psycopg2\n",
    "\n",
    "def connect_db():\n",
    "    return psycopg2.connect(DATABASE_CONNECTION_STRING)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the source table\n",
    "\n",
    "Before we create the source table, we ensure that the latest version (<u>**0.4.0**</u>) of [pgai](https://github.com/timescale/pgai?tab=readme-ov-file#pgai-allows-you-to-develop-rag-semantic-search-and-other-ai-applications-directly-in-postgresql) is installed in our database service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "with connect_db() as conn:\n",
    "    with conn.cursor() as cur:\n",
    "         cur.execute(\"CREATE EXTENSION IF NOT EXISTS ai CASCADE;\")\n",
    "\n",
    "    with conn.cursor() as cur:\n",
    "         cur.execute(\"\"\"\n",
    "                CREATE TABLE IF NOT EXISTS products(\n",
    "                    id bigint primary key generated by default as identity,\n",
    "                    product TEXT NOT NULL,\n",
    "                    description TEXT NOT NULL\n",
    "                );\n",
    "          \"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inserting `subset_1` in the source table\n",
    "\n",
    "We insert the data in `subset_1` to demonstrate how the vectorizer is ready to generate embeddings as soon as it configured and created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "\n",
    "## Helper functions\n",
    "def insert_subset_data(subset):\n",
    "    with connect_db() as conn:\n",
    "        with conn.cursor() as cur:\n",
    "            output = io.StringIO()\n",
    "\n",
    "            for product in subset:\n",
    "                output.write(f\"{product['name']}\\t{product['description']}\\n\")\n",
    "            output.seek(0)\n",
    "\n",
    "            cur.copy_from(output,'products', columns=('product', 'description'),\n",
    "                          sep='\\t')\n",
    "\n",
    "def count_embeddings():\n",
    "    with connect_db() as conn:\n",
    "        with conn.cursor() as cur:\n",
    "            cur.execute(\"SELECT COUNT(*) FROM products_embedding_store;\")\n",
    "            print(\"Number of vector embeddings generated:\", cur.fetchone()[0])\n",
    "\n",
    "def selecting_embeddings(table_name):\n",
    "    query = f\"SELECT * FROM {table_name} LIMIT 1;\"\n",
    "\n",
    "    with connect_db() as conn:\n",
    "        with conn.cursor() as cur:\n",
    "            cur.execute(query)\n",
    "            print(cur.fetchone())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "insert_subset_data(subset_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check that the insert operation was successful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with connect_db() as conn:\n",
    "    with conn.cursor() as cur:\n",
    "        cur.execute(\"SELECT COUNT(*) FROM products;\")\n",
    "        print(\"Number of items inserted:\", cur.fetchone()[0])\n",
    "\n",
    "        cur.execute(\"SELECT * FROM products LIMIT 1;\")\n",
    "        print(cur.fetchone())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the vectorizer\n",
    "\n",
    "In this tutorial, we use the cloud option for creating the vectorizer. Therefore:\n",
    "\n",
    "- navigate to [Timescale console](https://console.cloud.timescale.com/signup/?utm_source=blog&utm_medium=website&utm_campaign=vectorlaunch&utm_content=automate-embeddings-cta)\n",
    "- setup the OpenAI API Key in the credentials (specifically `AI Model API Keys`) of your [project's settings](https://api.dev.metronome-cloud.com/dashboard/settings).\n",
    "- then open your service and navigate to the `AI`tab where you can create a vectorizer. \n",
    "\n",
    "Here is the finalized SQL command, I use to create the vectorizer. It can look different depending on the parameters you choose to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "sql"
    }
   },
   "outputs": [],
   "source": [
    "SELECT ai.create_vectorizer(\n",
    "    'public.products'::regclass\n",
    "  , embedding=>ai.embedding_openai('text-embedding-3-small', 1536, api_key_name=>'OPENAI_API_KEY')\n",
    "  , chunking=>ai.chunking_recursive_character_text_splitter('description')\n",
    "  , formatting=>ai.formatting_python_template('product: $product description: $chunk')\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Embeddings table\n",
    "\n",
    "Once your vectorizer is up and running, you can query the embeddings table to see the embeddings generated. The naming convention for the embeddings table is `<source_table>_embedding_store`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selecting_embeddings('products_embedding_store')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View of source data & embeddings\n",
    "\n",
    "pgai vectorizer also creates a view of the source data and embeddings together. This view is named following the convention: `<source_table>_embedding`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selecting_embeddings('products_embedding')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieve context: Query the embeddings\n",
    "\n",
    "Let's perform vector similarity search on the embeddings generated by the vectorizer to retrieve context for our query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"Tell me about different types of t-shirts.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with connect_db() as conn:\n",
    "    with conn.cursor() as cur:\n",
    "        cur.execute(\"\"\"\n",
    "            SELECT chunk\n",
    "            FROM products_embedding\n",
    "            ORDER BY embedding <=> ai.openai_embed(\n",
    "                'text-embedding-3-small', %s\n",
    "            ) ASC\n",
    "            LIMIT 3;\n",
    "        \"\"\", (query,))\n",
    "\n",
    "        rows = cur.fetchall()\n",
    "\n",
    "        # Prepare the context for generating the response\n",
    "        context = \"\\n\\n\".join([f\"Chunk: {row[0]}\" for row in rows])\n",
    "        print(context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate model response\n",
    "\n",
    "Use [`GPT-4o mini`](https://platform.openai.com/docs/models/gpt-4o-mini) from OpenAI to generate a response from the query and context defined in the previous section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with connect_db() as conn:\n",
    "    with conn.cursor() as cur:\n",
    "        cur.execute(\n",
    "            \"\"\"\n",
    "            SELECT ai.openai_chat_complete(\n",
    "                'gpt-4o-mini',\n",
    "                jsonb_build_array(\n",
    "                    jsonb_build_object('role', 'system', 'content', \n",
    "                                       'you are a helpful assistant'),\n",
    "                    jsonb_build_object('role', 'user', 'content', %s)\n",
    "                )\n",
    "            )::json->'choices'->0->'message'->>'content';\n",
    "            \"\"\",\n",
    "            (f\"Query: {query}\\nContext: {context}\",)\n",
    "        )\n",
    "\n",
    "        # Fetch the response from the model\n",
    "        model_response = cur.fetchone()[0]\n",
    "        print(model_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Updating embeddings: Inserting new data in the source table\n",
    "\n",
    "We insert `subset_2` in our source data table: `products`. As soon as the new items are inserted, the vectorizer schedules their processing and starts running shortly after eventually being up to date all without any other code written."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_embeddings() ## before updating the source data.\n",
    "insert_subset_data(subset_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can go check the status of your vectorizer on [Timescale console](https://console.cloud.timescale.com/signup/?utm_source=blog&utm_medium=website&utm_campaign=vectorlaunch&utm_content=automate-embeddings-cta). Once the vectorizer indicates that it is updated, you can run the code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_embeddings() ## before updating the source data."
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
