{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "id": "iyi0u1inBcr1" }, "outputs": [], "source": [ "# Copyright 2024 Google LLC\n", "#\n", "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", "# you may not use this file except in compliance with the License.\n", "# You may obtain a copy of the License at\n", "#\n", "# https://www.apache.org/licenses/LICENSE-2.0\n", "#\n", "# Unless required by applicable law or agreed to in writing, software\n", "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "# See the License for the specific language governing permissions and\n", "# limitations under the License." ] }, { "cell_type": "markdown", "metadata": { "id": "BLIlkXRfCNyQ" }, "source": [ "# Retrieval Augmented Generation(RAG) with AlloyDB\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", "
\n", " \n", " \"Google
Open in Colab\n", " \n", " 		\n", " \n", " \"Google
Open in Colab Enterprise\n", " \n", " 		\n", " \n", " \"Vertex
Open in Vertex AI Workbench\n", " \n", " 		\n", " \n", " \"BigQuery
Open in BigQuery Studio\n", " \n", " 		\n", " \n", " \"GitHub
View on GitHub\n", " \n", "
\n", "\n", "
\n", "\n", "Share to:\n", "\n", "\n", " \"LinkedIn\n", "\n", "\n", "\n", " \"Bluesky\n", "\n", "\n", "\n", " \"X\n", "\n", "\n", "\n", " \"Reddit\n", "\n", "\n", "\n", " \"Facebook\n", " " ] }, { "cell_type": "markdown", "metadata": { "id": "WBpYfJP8Ce4s" }, "source": [ "| | |\n", "|--------------------------------------------|---------------------------------------------------|\n", "|Author(s) | |\n", "|[Tanya Warrier](https://github.com/tanyarw) |[Rupjit Chakraborty](https://github.com/lazyprgmr) |" ] }, { "cell_type": "markdown", "metadata": { "id": "WZPabwVQCh1y" }, "source": [ "## Overview\n", "\n", "- **PostgreSQL:** [PostgreSQL](https://www.postgresql.org/docs/current/) is an open-source, highly-extensible object-relational database management system known for its reliability and feature richness.\n", "\n", "- **AlloyDB:** [AlloyDB](https://cloud.google.com/alloydb/docs/overview) is Google Cloud's fully-managed, PostgreSQL-compatible database service optimized for demanding enterprise workloads and transactional/analytical hybrid processing.\n", "\n", "- **Gemini:** [Gemini](https://ai.google.dev/models/gemini) is a family of generative AI models that lets developers generate content and solve problems. These models are designed and trained to handle both text and images as input.\n", "\n", "- **Vertex AI Embeddings for Text:** With [text-embedding-005](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings) models we can easily create a text embedding with LLM. `text-embedding-005` is the newest stable embedding model.\n", "\n", "This notebook demonstrates Retrieval Augmented Generation (RAG) with AlloyDB backend. After installing the pre-requisites, we create an AlloyDB instance and use it to store embeddings. Finally we demonstrate how to fetch similar documents from AlloyDB and answer questions based on the documents fetched using Gemini.\n", " \n", "Text embeddings are created for publicly available abstracts from patents data and use them in our LLM search. Google Patents Research Data contains the output of much of the data analysis work used in Google Patents (patents.google.com).\n", " \n", "**Dataset**: `patents-public-data.google_patents_research.publications`\n" ] }, { "cell_type": "markdown", "metadata": { "id": "fdU7KNcCCkOA" }, "source": [ "## Getting Started" ] }, { "cell_type": "markdown", "metadata": { "id": "skNFqrfxuT1f" }, "source": [ "### Enable Cloud APIs\n", "Google Cloud APIs are programmatic interfaces to Google Cloud Platform services.\n", "\n", "1. [Recommended APIs for AlloyDB](https://cloud.google.com/alloydb/docs/project-enable-access)\n", "2. [Recommended APIs for Vertex AI](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis)\n", "\n", "\n", "#### **Before Moving Forward**\n", "\n", "Ensure your project has [private services access](https://cloud.google.com/alloydb/docs/about-private-services-access) enabled with `Google Cloud Platform` as the service provider.\n", "\n", "For instructions to set it up, click [here](https://cloud.google.com/alloydb/docs/configure-connectivity)." ] }, { "cell_type": "markdown", "metadata": { "id": "6XczCHZYCmi1" }, "source": [ "### Install required packages\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "4yzI1lAq9lPf" }, "outputs": [], "source": [ "%pip install pg8000==1.31.1 \\\n", "SQLAlchemy==2.0.29 \\\n", "google-cloud-aiplatform==1.46.0 \\\n", "google-cloud-alloydb-connector==1.0.0 --upgrade --quiet" ] }, { "cell_type": "markdown", "metadata": { "id": "bHV6vh55Cqln" }, "source": [ "### Restart runtime (Colab only)\n", "\n", "To use the newly installed packages, you must restart the runtime on Google Colab." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "mB85g9FpDu0e" }, "outputs": [], "source": [ "import sys\n", "\n", "if \"google.colab\" in sys.modules:\n", " import IPython\n", "\n", " app = IPython.Application.instance()\n", " app.kernel.do_shutdown(True)" ] }, { "cell_type": "markdown", "metadata": { "id": "VA-peJUWCujf" }, "source": [ "
\n", "⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️\n", "
" ] }, { "cell_type": "markdown", "metadata": { "id": "K810s6DICwe1" }, "source": [ "### Authenticate your notebook environment (Colab only)\n", "\n", "\n", "If you are running this notebook on Google Colab, run the cell below to authenticate your environment." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "K0xFZPSpDcog" }, "outputs": [], "source": [ "import sys\n", "\n", "if \"google.colab\" in sys.modules:\n", " from google.colab import auth\n", "\n", " auth.authenticate_user()" ] }, { "cell_type": "markdown", "metadata": { "id": "ooauhzOWPvC1" }, "source": [ "### Import Libraries\n", "\n", "Imports and prepares libraries for interacting with AlloyDB, Vertex AI resources for models, and data manipulation.\n", "\n", "* `pg8000:` PostgreSQL database driver.\n", "* `vertexai:` Google Cloud Vertex AI for managing generative models.\n", "* `sqlalchemy:` ORM for interfacing with databases in Python.\n", "* `subprocess:` Spawns external processes.\n", "* `pandas:` Powerful library for data analysis and manipulation.\n", "* `google.cloud.alloydb.connector:` Connector for Google Cloud AlloyDB (managed PostgreSQL)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Hj6K6SFIEnRk" }, "outputs": [], "source": [ "import subprocess\n", "\n", "from google.cloud.alloydb.connector import Connector\n", "import pandas as pd\n", "import pg8000\n", "import sqlalchemy\n", "from sqlalchemy.engine import Engine\n", "from sqlalchemy.exc import DatabaseError\n", "import vertexai\n", "from vertexai.generative_models import GenerationConfig, GenerativeModel" ] }, { "cell_type": "markdown", "metadata": { "id": "nPtpzZd7zU_e" }, "source": [ "### Set Google Cloud project information and initialize Vertex AI SDK\n", "\n", "- To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n", "\n", "- Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "5OwGCuGVvxFQ" }, "outputs": [], "source": [ "PROJECT_ID = \"\" # @param {type:\"string\"}\n", "LOCATION = \"us-central1\" # @param {type:\"string\"}\n", "\n", "vertexai.init(project=PROJECT_ID, location=LOCATION)" ] }, { "cell_type": "markdown", "metadata": { "id": "0C2kQjTmzYL9" }, "source": [ "Configurations to create a new AlloyDB cluster and primary instance\n", "\n", "* To get started using AlloyDB, you must have an existing Google Cloud project and [enable the AlloyDB AI API](https://console.cloud.google.com/flows/enableapi?apiid=alloydb.googleapis.com).\n", "* To generate embeddings with AlloyDB, the cluster created must reside in the region `us-central1`.\n", "* Please follow the naming convention listed below for `CLUSTER` and `INSTANCE` name\n", " - Must begin with a lowercase letter.\n", " - Can optionally contain a combination of lowercase letters, numbers, and hyphens (up to 61 characters).\n", " - Must end with a lowercase letter or digit." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "mv3ho2AwDkID" }, "outputs": [], "source": [ "REGION = \"us-central1\" # @param {type:\"string\"}\n", "CLUSTER = \"\" # @param {type:\"string\"}\n", "INSTANCE = \"\" # @param {type:\"string\"}\n", "CPU_COUNT = 2 # @param {type:\"integer\"}\n", "\n", "PROJECT_NUM = (\n", " subprocess.check_output(\n", " [\n", " \"gcloud\",\n", " \"projects\",\n", " \"describe\",\n", " PROJECT_ID,\n", " \"--format\",\n", " \"value(projectNumber)\",\n", " ]\n", " )\n", " .decode(\"utf-8\")\n", " .strip()\n", ")\n", "print(f\"Project Number: {PROJECT_NUM}\")\n", "\n", "SERVICE_ACCOUNT = (\n", " f\"serviceAccount:service-{PROJECT_NUM}@gcp-sa-alloydb.iam.gserviceaccount.com\"\n", ")\n", "print(f\"AlloyDB Service Agent: {SERVICE_ACCOUNT}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "hV7-Maoczm9W" }, "source": [ "Configurations for embedding and generative model\n", "\n", "\n", "* `text-embedding-005` outputs 768-dimensional vector embeddings.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "nLjwxgWvxJ0q" }, "outputs": [], "source": [ "EMBEDDING_MODEL = \"text-embedding-005\" # @param {type:\"string\"}\n", "DIMENSIONS = 768 # @param {type:\"integer\"}\n", "GENERATIVE_MODEL = \"gemini-2.0-flash\" # @param {type:\"string\"}" ] }, { "cell_type": "markdown", "metadata": { "id": "8rxwu-TaSdmn" }, "source": [ "## Fetch Dataset from BigQuery\n", "\n", "- [Google Patents Research Data](https://console.cloud.google.com/marketplace/product/google_patents_public_datasets/google-patents-research-data) contains the output of much of the data analysis work used in [Google Patents](https://patents.google.com), including machine translations of titles and abstracts from Google Translate, embedding vectors, extracted top terms, similar documents, and forward references.\n", "- We will use the public dataset table `google_patents_research.publications` for this demo by selecting the text columns below:\n", " - `publication_number:` Patent publication number (DOCDB compatible), eg: 'US-7650331-B1'.\n", " - `title:` The English title.\n", " - `abstract:` The English abstract.\n", " - `url:` URL to the patents.google.com page for this patent.\n", " - `country:` Country name.\n", " - `publication_description:` Description of the publication type.\n", "- The text columns `title` and `abstract` will later be converted into text embeddings to perform similarity search. The other columns would be used as supplemental information to the user's question." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "rhfMamgG6kZf" }, "outputs": [], "source": [ "query = \"\"\"\n", "SELECT publication_number,\ttitle,\tabstract, url,\tcountry,\tpublication_description\n", "FROM `patents-public-data.google_patents_research.publications`\n", "WHERE\n", " length(title)>1\n", " AND\n", " length(abstract)>1\n", "ORDER BY \tpublication_number\n", "LIMIT 1000\n", "\"\"\"\n", "# Read the table and display first 5 rows\n", "df = pd.read_gbq(query, project_id=PROJECT_ID)\n", "df.head(5)" ] }, { "cell_type": "markdown", "metadata": { "id": "zTCmZN_H7fHx" }, "source": [ "## AlloyDB as RAG backend" ] }, { "cell_type": "markdown", "metadata": { "id": "3hqFn7KPVAtX" }, "source": [ "### Set Up\n", "\n", "Set up cluster, instance and update them to allow public IP.\n", " - First we must create and connect the database on AlloyDB for\n", "PostgreSQL. For more details check [create a cluster](https://cloud.google.com/alloydb/docs/cluster-create) and [create a primary instance](https://cloud.google.com/alloydb/docs/instance-primary-create).\n", " - To generate embeddings with AlloyDB, the created cluster must reside in the region `us-central1`. This is required because the Vertex AI model that AlloyDB can use for embeddings, `text-embedding-005`, is located in that region.\n", "More details about embedding generation can be found [here](https://codelabs.developers.google.com/codelabs/alloydb-ai-embedding)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Y7CfxBOZLQ2A" }, "outputs": [], "source": [ "password = input(\"Enter a password for the cluster: \")\n", "\n", "# Set the active Google Cloud Project\n", "!gcloud config set project {PROJECT_ID}\n", "\n", "# Create cluster\n", "!gcloud alloydb clusters create {CLUSTER} --password={password} --region={REGION}\n", "\n", "# Create the primary instance\n", "!gcloud alloydb instances create {INSTANCE} --instance-type=PRIMARY --cpu-count={CPU_COUNT} --region={REGION} --cluster={CLUSTER}\n", "\n", "# Update the instance to allow public IP\n", "!gcloud beta alloydb instances update {INSTANCE} --region={REGION} --cluster={CLUSTER} --assign-inbound-public-ip=ASSIGN_IPV4" ] }, { "cell_type": "markdown", "metadata": { "id": "zV1B1C5NCOvF" }, "source": [ "Provision the `aiplatform.user` role to the AlloyDB service agent" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "774zMpFyowgF" }, "outputs": [], "source": [ "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n", " --member={SERVICE_ACCOUNT} \\\n", " --role=\"roles/aiplatform.user\"" ] }, { "cell_type": "markdown", "metadata": { "id": "uaaO4y-1gDYd" }, "source": [ "
\n", "⚠️ Please wait for a few minutes to ensure that the AlloyDB instance is updated with a Public IP address before moving forward. ⚠️\n", "
" ] }, { "cell_type": "markdown", "metadata": { "id": "L8iee-ydck4K" }, "source": [ "### Helper Functions\n", "\n", "\n", "* `create_sqlalchemy_engine`: Creates connection pool for AlloyDB instance\n", "* `check_table_exists`: Checks if table exists in an instance\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "0scMfoar7d4B" }, "outputs": [], "source": [ "def create_sqlalchemy_engine(\n", " inst_uri: str, user: str, password: str, db: str\n", ") -> tuple[sqlalchemy.engine.Engine, Connector]:\n", " \"\"\"Creates a connection pool for an AlloyDB instance and returns the pool\n", " and the connector. Callers are responsible for closing the pool and the\n", " connector.\n", "\n", "\n", " Args:\n", " inst_uri (str):\n", " The instance URI specifies the instance relative to the project,\n", " region, and cluster. For example:\n", " \"projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance\"\n", " user (str):\n", " The database user name, e.g., postgres\n", " password (str):\n", " The database user's password, e.g., secret-password\n", " db (str):\n", " The name of the database, e.g., mydb\n", "\n", " Returns:\n", " Tuple[sqlalchemy.engine.Engine, Connector]:\n", " * A SQLAlchemy engine object for managing database interactions.\n", " * A Connector object for underlying database connections (can be used for closing).\n", " \"\"\"\n", " connector = Connector()\n", "\n", " def getconn() -> pg8000.dbapi.Connection:\n", " \"\"\"\n", " Establishes a connection to a Google Cloud AlloyDB instance (PostgreSQL database) using the pg8000 driver.\n", "\n", " Returns:\n", " pg8000.dbapi.Connection: An active database connection object.\n", " \"\"\"\n", " conn: pg8000.dbapi.Connection = connector.connect(\n", " instance_uri=inst_uri,\n", " driver=\"pg8000\",\n", " user=user,\n", " password=password,\n", " db=db,\n", " ip_type=\"PUBLIC\", # use ip_type to specify Public IP\n", " )\n", " return conn\n", "\n", " # create SQLAlchemy connection pool\n", " engine = sqlalchemy.create_engine(\n", " \"postgresql+pg8000://\", creator=getconn, isolation_level=\"AUTOCOMMIT\"\n", " )\n", " engine.dialect.description_encoding = None\n", " return engine, connector" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "bEv-Yj5MTBoM" }, "outputs": [], "source": [ "def check_table_exists(engine: Engine, connector: Connector, table_name: str) -> bool:\n", " \"\"\"Checks if a table exists in the database.\n", "\n", " Args:\n", " engine (sqlalchemy.engine.Engine): SQLAlchemy engine object.\n", " connector (Connector): AlloyDB Connector object.\n", " table_name (str): Name of the table to check.\n", "\n", " Returns:\n", " bool: True if the table exists, False otherwise.\n", " \"\"\"\n", "\n", " try:\n", " with engine.connect() as conn:\n", " check_cmd = sqlalchemy.text(f\"SELECT 1 FROM {table_name} LIMIT 1\")\n", " conn.execute(check_cmd)\n", " connector.close()\n", " return True\n", "\n", " except DatabaseError:\n", " return False" ] }, { "cell_type": "markdown", "metadata": { "id": "1hAafWeecyeG" }, "source": [ "### Create the connection to AlloyDB" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "F3DnzTH3LRz5" }, "outputs": [], "source": [ "INSTANCE_URI = (\n", " f\"projects/{PROJECT_ID}/locations/{REGION}/clusters/{CLUSTER}/instances/{INSTANCE}\"\n", ")\n", "USER = \"postgres\"\n", "DB = \"postgres\"\n", "TABLE_NAME = \"google_patents_research\"" ] }, { "cell_type": "markdown", "metadata": { "id": "TcfvIbVbkamU" }, "source": [ "### Create a table on AlloyDB\n", "The table is created with the columns from the `google_patents_research.publications` dataset.\n", "\n", "> **Note:** If you come across the error below, it is because the AlloyDB instance has not finished updating its public IP address. Please wait for a few minutes before trying to assign it again under the **Set Up** section.\n", "```\n", "IPTypeNotFoundError: AlloyDB instance does not have an IP addresses matching type: 'PUBLIC'\n", "```" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "VUaDKZnWDDo2" }, "outputs": [], "source": [ "engine, connector = create_sqlalchemy_engine(\n", " inst_uri=INSTANCE_URI,\n", " user=USER,\n", " password=password,\n", " db=DB,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Xq8qgcewevsY" }, "outputs": [], "source": [ "if check_table_exists(engine, connector, TABLE_NAME):\n", " print(f\"Table {TABLE_NAME} already exists!\")\n", "\n", "else:\n", " # Create table\n", " create_table_cmd = sqlalchemy.text(\n", " f\"CREATE TABLE {TABLE_NAME} ( \\\n", " publication_number VARCHAR, \\\n", " title TEXT, \\\n", " abstract TEXT, \\\n", " url VARCHAR, \\\n", " country TEXT, \\\n", " publication_description TEXT \\\n", " )\",\n", " )\n", "\n", " # Insert data\n", " insert_data_cmd = sqlalchemy.text(\n", " f\"\"\"\n", " INSERT INTO {TABLE_NAME} VALUES (:publication_number, :title,\t:abstract, :url,\t:country,\t:publication_description)\n", " \"\"\"\n", " )\n", "\n", " parameter_map = [\n", " {\n", " \"publication_number\": row[\"publication_number\"],\n", " \"title\": row[\"title\"],\n", " \"abstract\": row[\"abstract\"],\n", " \"url\": row[\"url\"],\n", " \"country\": row[\"country\"],\n", " \"publication_description\": row[\"publication_description\"],\n", " }\n", " for index, row in df.iterrows()\n", " ]\n", "\n", " # Execute the queries\n", " with engine.connect() as conn:\n", " print(\"Creating table...\")\n", " conn.execute(create_table_cmd)\n", " print(\"Inserting values...\")\n", " conn.execute(\n", " insert_data_cmd,\n", " parameter_map,\n", " )\n", " print(\"Committing...\")\n", " conn.commit()\n", " print(\"Done\")\n", " connector.close()" ] }, { "cell_type": "markdown", "metadata": { "id": "l_vCJCV_c5ym" }, "source": [ "### Add AlloyDB extensions\n", "\n", "Enable an extension by connecting to a database in an AlloyDB cluster's primary instance, then running a `CREATE EXTENSION` command. More details can be found [here](https://cloud.google.com/alloydb/docs/reference/extensions#enable).\n", "- `google_ml_integration` integrates AlloyDB with Vertex AI\n", "- `vector` allows us to use `pgvector` functions and operators with optimizations specific to AlloyDB" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "lXeVfDinBxMj" }, "outputs": [], "source": [ "# Add extensions\n", "google_ml_integration_cmd = sqlalchemy.text(\n", " \"CREATE EXTENSION IF NOT EXISTS google_ml_integration CASCADE\"\n", ")\n", "vector_cmd = sqlalchemy.text(\"CREATE EXTENSION IF NOT EXISTS vector\")\n", "\n", "# Execute the queries\n", "with engine.connect() as conn:\n", " conn.execute(google_ml_integration_cmd)\n", " conn.execute(vector_cmd)\n", " conn.commit()\n", "connector.close()" ] }, { "cell_type": "markdown", "metadata": { "id": "2gmhxd4Nkgpv" }, "source": [ "### Create a column that stores text embeddings and an Index using AlloyDB\n", "\n", "\n", "\n", "* The Vertex AI text-embeddings API lets you create a text embedding using Generative AI on Vertex AI. Text embeddings are numerical representations of text that capture relationships between words and phrases.\n", "* IVFFlat is a type of vector index for approximate nearest neighbor search. It is a frequently used index type that can improve performance when querying highly-dimensional vectors, like those representing embeddings.\n", "\n", "Visit the [pgvector documentation](https://github.com/pgvector/pgvector?tab=readme-ov-file#pgvector) for more information on supported index types and their distance functions\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "nQ2Xw15HgQ7L" }, "outputs": [], "source": [ "embedding_column = \"embedding\"\n", "distance_function = \"vector_cosine_ops\"\n", "\n", "# Add column to store embeddings\n", "add_column_cmd = sqlalchemy.text(\n", " f\"ALTER TABLE {TABLE_NAME} ADD COLUMN {embedding_column} vector({DIMENSIONS});\"\n", ")\n", "\n", "# Generate embeddings for `title` and `abstract` columns of the dataset\n", "embedding_cmd = sqlalchemy.text(\n", " f\"UPDATE {TABLE_NAME} SET {embedding_column} = embedding('{EMBEDDING_MODEL}', title || ' ' || abstract);\"\n", ")\n", "\n", "# Create an IVFFlat index on the table with embedding column and cosine distance\n", "index_cmd = sqlalchemy.text(\n", " f\"CREATE INDEX ON {TABLE_NAME} USING ivfflat ({embedding_column} {distance_function})\"\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "jfBtZUI8mQJU" }, "outputs": [], "source": [ "# Execute the queries\n", "with engine.connect() as conn:\n", " try:\n", " conn.execute(add_column_cmd)\n", " except:\n", " print(f\"Column {embedding_column} already exists\")\n", " print(\"Creating Embeddings...\")\n", " conn.execute(embedding_cmd)\n", " print(\"Creating Index...\")\n", " conn.execute(index_cmd)\n", " print(\"Committing...\")\n", " conn.commit()\n", " print(\"Done\")\n", "connector.close()" ] }, { "cell_type": "markdown", "metadata": { "id": "PwYsLCD9c_MJ" }, "source": [ "## Retrieve data\n", "\n", "Retrieve top 5 rows based on similarity search" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "5pSbkfBbwDL1" }, "outputs": [], "source": [ "def retrieve_information(\n", " query: str,\n", " engine: Engine,\n", " table_name: str,\n", " embedding_model: str,\n", " row_count: int = 5,\n", ") -> str:\n", " \"\"\"\n", " Queries a database table using a semantic similarity search and returns formatted results.\n", "\n", " Args:\n", " query (str): The search query to embed and compare against the database.\n", " engine (sqlalchemy.engine.Engine): SQLAlchemy engine object.\n", " table_name (str): The name of the table to query.\n", " embedding_model (str): The name of the embedding model to use.\n", " row_count (int, optional): The maximum number of results to return. Defaults to 5.\n", "\n", " Assumptions:\n", " The table has columns named 'publication_number', 'title', 'abstract', 'url', and an embedding column named 'embedding_column'.\n", "\n", " Returns:\n", " str: A formatted string containing the top results, including their publication number, title, abstract, and URL.\n", " \"\"\"\n", "\n", " # Perform semantic search\n", " search_cmd = sqlalchemy.text(\n", " f\"\"\"\n", " SELECT publication_number, title,\tabstract, url FROM {table_name}\n", " ORDER BY {embedding_column}\n", " <-> embedding('{embedding_model}', '{query}')::vector\n", " LIMIT {row_count}\n", " \"\"\"\n", " )\n", "\n", " # Execute the query\n", " with engine.connect() as conn:\n", " result = conn.execute(search_cmd)\n", " context = [row._asdict() for row in result]\n", " connector.close()\n", "\n", " # String format the retrieved information\n", " retrieved_information = \"\\n\".join(\n", " [\n", " f\"{index+1}. \"\n", " + \"\\n\".join([f\"{key}: {value}\" for key, value in element.items()])\n", " for index, element in enumerate(context)\n", " ]\n", " )\n", "\n", " return retrieved_information" ] }, { "cell_type": "markdown", "metadata": { "id": "E7cjd66wHaUx" }, "source": [ "**Sample Questions**\n", "\n", "\n", "1. Propose some project ideas for medical devices.\n", "2. List patents around solar energy and how can they be used.\n", "3. What methods exist to improve combustion?\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "vrDZ8hB6FjkX" }, "outputs": [], "source": [ "query = \"Propose some project ideas for medical devices.\" # @param {type:\"string\"}\n", "\n", "result = retrieve_information(\n", " query=query, engine=engine, table_name=TABLE_NAME, embedding_model=EMBEDDING_MODEL\n", ")\n", "print(result)" ] }, { "cell_type": "markdown", "metadata": { "id": "8img1AVdqdtB" }, "source": [ "## Generate Response" ] }, { "cell_type": "markdown", "metadata": { "id": "ubQzJ-gYKdLr" }, "source": [ "Define a prompt template to answer questions according to the use-case." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "txb2LUMaxSCm" }, "outputs": [], "source": [ "prompt = \"\"\"You are a friendly advisor helping to answer questions about patents. Based on the search request we have loaded a list of patents closely related to the search.\n", "\n", "The user asked:\n", "\n", "{question}\n", "\n", "\n", "Here is the list of matching patents:\n", "\n", "{result}\n", "\n", "\n", "You should answer the question using the matching patents, reply with supplemental information and patent url.\n", "Answer:\n", "\"\"\"" ] }, { "cell_type": "markdown", "metadata": { "id": "gcgf_RruKrJC" }, "source": [ "The `generate_text` function performs two tasks\n", "- Formats the prompt template with `result` and `question`.\n", "- Invokes the generative model, in this case Gemini." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "5Elqguzq-SHY" }, "outputs": [], "source": [ "def generate_text(\n", " prompt: str,\n", " result: str,\n", " question: str,\n", " generative_model: GenerativeModel,\n", " generation_config: GenerationConfig,\n", ") -> str:\n", " \"\"\"\n", " Generates text response using a specified generative language model on Vertex AI.\n", "\n", " Args:\n", " prompt (str): The text prompt template for the generative model.\n", " result (str): The list of matching patents.\n", " question (str): The user's question.\n", " generative_model (vertexai.generative_models.GenerativeModel): The name or identifier of the generative model on Vertex AI.\n", " generation_config (vertexai.generative_models.GenerationConfig): Configuration object for the text generation process.\n", "\n", " Returns:\n", " str: The generated text response from the model.\n", " \"\"\"\n", " input_prompt = prompt.format(result=result, question=question)\n", "\n", " # Query the model\n", " response = generative_model.generate_content(\n", " contents=input_prompt, generation_config=generation_config\n", " )\n", "\n", " return response.text" ] }, { "cell_type": "markdown", "metadata": { "id": "MublwjbdKV7I" }, "source": [ "Generate a response based on the top 5 rows fetched via similarity search." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "xCanM1asBSXc" }, "outputs": [], "source": [ "response = generate_text(\n", " prompt=prompt,\n", " result=result,\n", " question=query,\n", " generative_model=GenerativeModel(GENERATIVE_MODEL),\n", " generation_config=GenerationConfig(temperature=0.6, max_output_tokens=1024),\n", ")\n", "\n", "print(response)" ] }, { "cell_type": "markdown", "metadata": { "id": "QXvQ_mThNxZZ" }, "source": [ "## Summary\n", "\n", "The provided code demonstrates a comprehensive approach to leveraging AlloyDB as a backend for a retrieval-augmented generative (RAG) application. The code is relevant to the task of building a RAG application using Vertex AI.\n", "\n", "AlloyDB delivers up to 100X faster analytical queries than standard PostgreSQL, and AlloyDB AI runs vector queries up to 10x faster compared to standard PostgreSQL when using the IVFFlat index.\n", "\n", "**Steps to improve generated responses:**\n", "\n", "1. Prompt Engineering: The prompt template used for text generation could be further refined to improve the relevance and quality of the responses based on the use-case.\n", "2. Contextual Information: Incorporate additional context from the retrieved information to provide more comprehensive and nuanced responses.\n", "3. Model Parameters: Explore and tune the parameters (`GenerationConfig`) of the model to enhance the quality and relevance of the generated responses." ] }, { "cell_type": "markdown", "metadata": { "id": "Th-slu0VCdDf" }, "source": [ "## Cleaning Up" ] }, { "cell_type": "markdown", "metadata": { "id": "ncKTgQdu2OPg" }, "source": [ "Clean up the created resources by deleting the primary instance, and the cluster." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "7LcIq5UevulE" }, "outputs": [], "source": [ "# Delete the instance\n", "!gcloud alloydb instances delete {INSTANCE} --cluster={CLUSTER} --region={REGION}\n", "\n", "# Delete the cluster\n", "!gcloud alloydb clusters delete {CLUSTER} --region={REGION}" ] } ], "metadata": { "colab": { "name": "rag_embeddings_and_index_with_alloydb.ipynb", "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }