Creating notebook to ingest CloudSQL database using kubernetes docs

applications/rag/README.md

@@ -17,7 +17,7 @@ RAG uses a semantically searchable knowledge base (like vector search) to retrie
 5. A [Jupyter](https://docs.jupyter.org/en/latest/) notebook running on GKE that reads the dataset using GCS fuse driver integrations and runs a Ray job to populate the vector DB.
 3. A front end chat interface running on GKE that prompts the inference server with context from the vector DB.
 
-This tutorial walks you through installing the RAG infrastructure in a GCP project, generating vector embeddings for a sample [Kaggle Netflix shows](https://www.kaggle.com/datasets/shivamb/netflix-shows) dataset and prompting the LLM with context.
+This tutorial walks you through installing the RAG infrastructure in a GCP project, generating vector embeddings for a sample [Kubernetes Docs](https://github.com/dohsimpson/kubernetes-doc-pdf) dataset and prompting the LLM with context.
 
 # Prerequisites
 
@@ -74,7 +74,7 @@ This section sets up the RAG infrastructure in your GCP project using Terraform.
 
 # Generate vector embeddings for the dataset
 
-This section generates the vector embeddings for your input dataset. Currently, the default dataset is [Netflix shows](https://www.kaggle.com/datasets/shivamb/netflix-shows). We will use a Jupyter notebook to run a Ray job that generates the embeddings & populates them into the `pgvector` instance created above.
+This section generates the vector embeddings for your input dataset. Currently, the default dataset is [Kubernetes docs](https://github.com/dohsimpson/kubernetes-doc-pdf). We will use a Jupyter notebook to generate the embeddings & populates them into the `pgvector` instance created above.
 
 Set your the namespace, cluster name and location from `workloads.tfvars`):
 
@@ -108,30 +108,10 @@ gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${CLUSTER_L
 
 2. Load the notebook:
     - Once logged in to JupyterHub, choose the `CPU` preset with `Default` storage. 
-    - Click [File] -> [Open From URL] and paste: `https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/main/applications/rag/example_notebooks/rag-kaggle-ray-sql-interactive.ipynb`
-
-3. Configure Kaggle:
-    - Create a [Kaggle account](https://www.kaggle.com/account/login?phase=startRegisterTab&returnUrl=%2F).
-    - [Generate an API token](https://www.kaggle.com/settings/account). See [further instructions](https://www.kaggle.com/docs/api#authentication). This token is used in the notebook to access the [Kaggle Netflix shows](https://www.kaggle.com/datasets/shivamb/netflix-shows) dataset.
-    - Replace the variables in the 1st cell of the notebook with your Kaggle credentials (can be found in the `kaggle.json` file created while generating the API token):
-        * `KAGGLE_USERNAME`
-        * `KAGGLE_KEY`
-
-4. Generate vector embeddings: Run all the cells in the notebook to generate vector embeddings for the Netflix shows dataset (https://www.kaggle.com/datasets/shivamb/netflix-shows) and store them in the `pgvector` CloudSQL instance via a Ray job.
-    * When the last cell says the job has succeeded (eg: `Job 'raysubmit_APungAw6TyB55qxk' succeeded`), the vector embeddings have been generated and we can launch the frontend chat interface. Note that running the job can take up to 10 minutes.
-    * Ray may take several minutes to create the runtime environment. During this time, the job will appear to be missing (e.g. `Status message: PENDING`).
-    * Connect to the Ray dashboard to check the job status or logs:
-        - If IAP is disabled (`ray_dashboard_add_auth = false`):
-            - `kubectl port-forward -n ${NAMESPACE} service/ray-cluster-kuberay-head-svc 8265:8265`
-            - Go to `localhost:8265` in a browser
-        - If IAP is enabled (`ray_dashboard_add_auth = true`):
-            - Fetch the domain: `terraform output ray-dashboard-managed-cert`
-            - If you used a custom domain, ensure you configured your DNS as described above.
-            - Verify the domain status is `Active`:
-                - `kubectl get managedcertificates ray-dashboard-managed-cert -n ${NAMESPACE} --output jsonpath='{.status.domainStatus[0].status}'`
-                - Note: This can take up to 20 minutes to propagate.
-            - Once the domain status is Active, go to the domain in a browser and login with your Google credentials.
-            - To add additional users to your frontend application, go to [Google Cloud Platform IAP](https://console.cloud.google.com/security/iap), select the `rag/ray-cluster-kuberay-head-svc` service and add principals with the role `IAP-secured Web App User`.
+    - Click [File] -> [Open From URL] and paste: `https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/main/applications/rag/example_notebooks/rag-data-ingest-with-kubernetes-docs.ipynb`
+
+
+3. Generate vector embeddings: Run all the cells in the notebook to generate vector embeddings for the [Kubernetes documentation](https://github.com/dohsimpson/kubernetes-doc-pdf) and store them in the `pgvector` CloudSQL instance using a Ray Job.
 
 # Launch the frontend chat interface
 

applications/rag/example_notebooks/rag-ray-ingest-with-kubernetes-docs.ipynb

@@ -0,0 +1,291 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "7e14d0f0-2573-4fe4-ba87-7a447f2f511c",
+   "metadata": {},
+   "source": [
+    "# RAG-on-GKE Application\n",
+    "\n",
+    "This is a Python notebook for generating the vector embeddings based on [Kubernetes docs](https://github.com/dohsimpson/kubernetes-doc-pdf/) used by the RAG on GKE application.   \n",
+    "For full information, please checkout the GitHub documentation [here](https://github.com/GoogleCloudPlatform/ai-on-gke/blob/main/applications/rag/README.md).\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2cba26cf",
+   "metadata": {},
+   "source": [
+    "## Clone the kubernetes docs repo"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5f9b1fad-537e-425f-a5fc-587a408b1fab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!mkdir /data/kubernetes-docs -p\n",
+    "!git clone https://github.com/dohsimpson/kubernetes-doc-pdf /data/kubernetes-docs\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b984429c-b65a-47b7-9723-ee3ad81d61db",
+   "metadata": {},
+   "source": [
+    "## Install the required packages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "40e4d29d-79c6-4233-a8ed-0f8a42576656",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "!pip install langchain langchain-community sentence_transformers pypdf"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f80cc5af-a1fa-456d-a4ed-fa2ffa3b87a0",
+   "metadata": {},
+   "source": [
+    "## Writting job to be used on the Ray Cluster"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "36523f3f-0c93-41da-abb9-c113bb456bc1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create a directory to package the contents that need to be downloaded in ray worker\n",
+    "! mkdir -p rag-app"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "69d912e5-2225-4b44-80cd-651f7cc71a40",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%writefile rag-app/job.py\n",
+    "\n",
+    "import os\n",
+    "import uuid\n",
+    "import glob\n",
+    "\n",
+    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
+    "from langchain.embeddings import HuggingFaceEmbeddings\n",
+    "from langchain_community.document_loaders import PyPDFLoader\n",
+    "\n",
+    "from google.cloud.sql.connector import Connector, IPTypes\n",
+    "import sqlalchemy\n",
+    "\n",
+    "from sqlalchemy.ext.declarative import declarative_base\n",
+    "from sqlalchemy import Column, String, Text, text, JSON\n",
+    "from sqlalchemy.orm import scoped_session, sessionmaker, mapped_column\n",
+    "from pgvector.sqlalchemy import Vector\n",
+    "\n",
+    "# initialize parameters\n",
+    "\n",
+    "INSTANCE_CONNECTION_NAME = os.environ[\"CLOUDSQL_INSTANCE_CONNECTION_NAME\"]\n",
+    "print(f\"Your instance connection name is: {INSTANCE_CONNECTION_NAME}\")\n",
+    "VECTOR_EMBEDDINGS_TABLE_NAME = \"rag_embeddings_db\"\n",
+    "DB_NAME = \"pgvector-database\"\n",
+    "\n",
+    "db_username_file = open(\"/etc/secret-volume/username\", \"r\")\n",
+    "DB_USER = db_username_file.read()\n",
+    "db_username_file.close()\n",
+    "\n",
+    "db_password_file = open(\"/etc/secret-volume/password\", \"r\")\n",
+    "DB_PASS = db_password_file.read()\n",
+    "db_password_file.close()\n",
+    "\n",
+    "# initialize Connector object\n",
+    "connector = Connector()\n",
+    "\n",
+    "# function to return the database connection object\n",
+    "def getconn():\n",
+    "    conn = connector.connect(\n",
+    "        INSTANCE_CONNECTION_NAME,\n",
+    "        \"pg8000\",\n",
+    "        user=DB_USER,\n",
+    "        password=DB_PASS,\n",
+    "        db=DB_NAME,\n",
+    "        ip_type=IPTypes.PRIVATE\n",
+    "    )\n",
+    "    return conn\n",
+    "\n",
+    "# create connection pool with 'creator' argument to our connection object function\n",
+    "pool = sqlalchemy.create_engine(\n",
+    "    \"postgresql+pg8000://\",\n",
+    "    creator=getconn,\n",
+    ")\n",
+    "\n",
+    "Base = declarative_base()\n",
+    "DBSession = scoped_session(sessionmaker())\n",
+    "\n",
+    "class TextEmbedding(Base):\n",
+    "    __tablename__ = VECTOR_EMBEDDINGS_TABLE_NAME\n",
+    "    id = Column(String(255), primary_key=True)\n",
+    "    text = Column(Text)\n",
+    "    text_embedding = mapped_column(Vector(384))\n",
+    "\n",
+    "with pool.connect() as conn:\n",
+    "    conn.execute(text(\"CREATE EXTENSION IF NOT EXISTS vector\"))\n",
+    "    conn.commit() \n",
+    "    \n",
+    "DBSession.configure(bind=pool, autoflush=False, expire_on_commit=False)\n",
+    "Base.metadata.drop_all(pool)\n",
+    "Base.metadata.create_all(pool)\n",
+    "\n",
+    "SENTENCE_TRANSFORMER_MODEL = \"intfloat/multilingual-e5-small\"  # Transformer to use for converting text chunks to vector embeddings\n",
+    "\n",
+    "# the dataset has been pre-dowloaded to the GCS bucket as part of the notebook in the cell above. Ray workers will find the dataset readily mounted.\n",
+    "SHARED_DATASET_BASE_PATH = \"/data/kubernetes-docs/\"\n",
+    "\n",
+    "CHUNK_SIZE = 1000  # text chunk sizes which will be converted to vector embeddings\n",
+    "CHUNK_OVERLAP = 10\n",
+    "VECTOR_DIMENSION = 384  # Embeddings size\n",
+    "\n",
+    "splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP, length_function=len)\n",
+    "embeddings_service = HuggingFaceEmbeddings(model_name=SENTENCE_TRANSFORMER_MODEL)\n",
+    "\n",
+    "def process_pdf(file_path):\n",
+    "    \"\"\"Loads, splits and embed a single PDF file.\"\"\"\n",
+    "    loader = PyPDFLoader(file_path)\n",
+    "    print(f\"Loading {file_path}\")\n",
+    "    pages = loader.load_and_split()\n",
+    "    \n",
+    "    splits = splitter.split_documents(pages)\n",
+    "\n",
+    "    chunks = []\n",
+    "    for split in splits:\n",
+    "        id = uuid.uuid4()\n",
+    "        page_content = split.page_content\n",
+    "        file_metadata = split.metadata\n",
+    "        embedded_document = embeddings_service.embed_query(page_content)\n",
+    "        split_data = {\n",
+    "            \"langchain_id\" : id,\n",
+    "            \"content\" : page_content,\n",
+    "            \"embedding\" : embedded_document,\n",
+    "            \"langchain_metadata\" : file_metadata\n",
+    "        }\n",
+    "        chunks.append(split_data)\n",
+    "    return chunks\n",
+    "\n",
+    "documents_file_paths = glob.glob(f\"{SHARED_DATASET_BASE_PATH}/PDFs/*.pdf\")\n",
+    "for file_path in documents_file_paths:\n",
+    "    processed_result = process_pdf(file_path)\n",
+    "    DBSession.bulk_insert_mappings(TextEmbedding, processed_result)\n",
+    "        \n",
+    "DBSession.commit()\n",
+    "\n",
+    "#Verifying the results.\n",
+    "\n",
+    "query_text = \"What's kubernetes?\" \n",
+    "query_emb = embeddings_service.embed_query(query_text).tolist()\n",
+    "query_request = \"SELECT id, text, text_embedding, 1 - ('[\" + \",\".join(map(str, query_emb)) + \"]' <=> text_embedding) AS cosine_similarity FROM \" + TABLE_NAME + \" ORDER BY cosine_similarity DESC LIMIT 5;\" \n",
+    "query_results = DBSession.execute(sqlalchemy.text(query_request)).fetchall()\n",
+    "DBSession.commit()\n",
+    "print(\"print query_results, the 1st one is the hit\")\n",
+    "for row in query_results:\n",
+    "    print(row)\n",
+    "\n",
+    "print (\"end job\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6b9bc582-50cd-4d7c-b5c4-549626fd2349",
+   "metadata": {},
+   "source": [
+    "## Summiting the job into Ray Cluster:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d5b6acbe-5a14-4bc8-a4ca-58a6b3dd5391",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import ray, time\n",
+    "from ray.job_submission import JobSubmissionClient\n",
+    "client = JobSubmissionClient(\"ray://ray-cluster-kuberay-head-svc:10001\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4eb8eae9-2a20-4c02-ac79-196942ae2783",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Port forward to the Ray dashboard and go to `localhost:8265` in a browser to see job status: kubectl port-forward -n <namespace> service/ray-cluster-kuberay-head-svc 8265:8265\n",
+    "import time\n",
+    "\n",
+    "start_time = time.time()\n",
+    "job_id = client.submit_job(\n",
+    "    entrypoint=\"python job.py\",\n",
+    "    # Path to the local directory that contains the entrypoint file.\n",
+    "    runtime_env={\n",
+    "        \"working_dir\": \"/home/jovyan/rag-app\", # upload the local working directory to ray workers\n",
+    "        \"pip\": [               \n",
+    "            \"langchain\",\n",
+    "            \"langchain-community\",\n",
+    "            \"sentence-transformers\",\n",
+    "            \"pypdf\",\n",
+    "            \"pgvector\"\n",
+    "        ]\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "# The Ray job typically takes 5m-10m to complete.\n",
+    "print(\"Job submitted with ID:\", job_id)\n",
+    "while True:\n",
+    "    status = client.get_job_status(job_id)\n",
+    "    print(\"Job status:\", status)\n",
+    "    print(\"Job info:\", client.get_job_info(job_id).message)\n",
+    "    if status.is_terminal():\n",
+    "        break\n",
+    "    time.sleep(30)\n",
+    "\n",
+    "end_time = time.time()\n",
+    "job_duration = end_time - start_time\n",
+    "print(f\"Job  completed in {job_duration} seconds.\")\n",
+    "\n",
+    "ray.shutdown()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

applications/rag/metadata.yaml

@@ -70,7 +70,7 @@ spec:
       - name: dataset_embeddings_table_name
         description: Name of the table that stores vector embeddings for input dataset
         varType: string
-        defaultValue: netflix_reviews_db
+        defaultValue: rag_embeddings_db
       - name: disable_ray_cluster_network_policy
         description: Disables Kubernetes Network Policy for Ray Clusters for this demo. Defaulting to 'true' aka disabled pending fixes to the kuberay-monitoring module. This should be defaulted to false.
         varType: bool