From a1d2738943c097a16dc041bba0fb5224fa8cf761 Mon Sep 17 00:00:00 2001 From: Alex Dorozhkin Date: Wed, 11 Dec 2024 15:43:45 -0500 Subject: [PATCH] feat: Notebook for RAG Engine with Vertex AI Search https://cloud.google.com/vertex-ai/generative-ai/docs/use-vertexai-search --- .../rag_engine_vertex_ai_search.ipynb | 670 ++++++++++++++++++ rag-grounding/README.md | 1 + 2 files changed, 671 insertions(+) create mode 100644 gemini/rag-engine/rag_engine_vertex_ai_search.ipynb diff --git a/gemini/rag-engine/rag_engine_vertex_ai_search.ipynb b/gemini/rag-engine/rag_engine_vertex_ai_search.ipynb new file mode 100644 index 0000000000..d291ec1fa7 --- /dev/null +++ b/gemini/rag-engine/rag_engine_vertex_ai_search.ipynb @@ -0,0 +1,670 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ur8xi4C7S06n" + }, + "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": "JAPoU8Sm5E6e" + }, + "source": [ + "# Vertex AI RAG Engine with Vertex AI Search\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", + " \"GitHub
View on GitHub\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "84f0f73a0f76" + }, + "source": [ + "| | |\n", + "|-|-|\n", + "| Author(s) | [Alex Dorozhkin](https://github.com/galexdor) |" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "tvgnzT1CKxrO" + }, + "source": [ + "## Overview\n", + "\n", + "This notebook illustrates how to use [Vertex AI RAG Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-overview) with [Vertex AI Search](https://cloud.google.com/enterprise-search) as a retrieval backend. Vertex AI Search's ability to handle large datasets, provide low-latency retrieval, and improve scalability makes it a powerful tool for enhancing RAG applications. By integrating Vertex AI Search, you can ensure that your RAG applications can efficiently access and process the necessary information for generating high-quality and contextually relevant responses.\n", + "\n", + "For more information, refer to the [official documentation](https://cloud.google.com/generative-ai-app-builder/docs/enterprise-search-introduction).\n", + "\n", + "For more details on RAG corpus/file management and detailed support please visit [Vertex AI RAG Engine API](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/rag-api)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "61RBz8LLbxCR" + }, + "source": [ + "## Get started" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "No17Cw5hgx12" + }, + "source": [ + "### Install Vertex AI SDK and other required packages\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "tFy3H3aPgx12" + }, + "outputs": [], + "source": [ + "%pip install --upgrade --user --quiet google-cloud-aiplatform google-cloud-discoveryengine" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "R5Xep4W9lq-Z" + }, + "source": [ + "### Restart runtime\n", + "\n", + "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n", + "\n", + "The restart might take a minute or longer. After it's restarted, continue to the next step." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "XRvKdaPDTznN" + }, + "outputs": [], + "source": [ + "import IPython\n", + "\n", + "app = IPython.Application.instance()\n", + "app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "SbmM4z7FOBpM" + }, + "source": [ + "
\n", + "⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️\n", + "
\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "DF4l8DTdWgPY" + }, + "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": 15, + "metadata": { + "id": "Nqwi-5ufWp_B" + }, + "outputs": [], + "source": [ + "# Use the environment variable if the user doesn't provide Project ID.\n", + "import os\n", + "\n", + "import vertexai\n", + "\n", + "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\", isTemplate: true}\n", + "if PROJECT_ID == \"[your-project-id]\":\n", + " PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n", + "\n", + "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n", + "\n", + "vertexai.init(project=PROJECT_ID, location=LOCATION)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "RQ-QWIZk6Rqb" + }, + "source": [ + "### Authenticate your notebook environment (Colab only)\n", + "\n", + "If you're running this notebook on Google Colab, run the cell below to authenticate your environment." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": { + "id": "FsX2KKvx7tm4" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + " from google.colab import auth\n", + "\n", + " auth.authenticate_user(project_id=PROJECT_ID)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "N4QCqfqJ36LR" + }, + "source": [ + "## (Optional) Setup Vertex AI Search Datastore and Engine" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "VxksV2zm6L1V" + }, + "source": [ + "In this section, we have some helper methods to help you setup your Vertex AI Search. These methods handle the creation of resources like Data Stores and Engines, which can take a few minutes.\n", + "\n", + "This section is not required if you already have a Vertex AI Search engine ready to use.\n", + "\n", + "To get started using Vertex AI Search, you must have an existing Google Cloud project and [enable the Discovery Engine API](https://console.cloud.google.com/flows/enableapi?apiid=discoveryengine.googleapis.com)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "XmvYhLurAta4" + }, + "source": [ + "### Initialize Vertex AI Search SDK" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "id": "UKX-J8OKAz3J" + }, + "outputs": [], + "source": [ + "from google.api_core.client_options import ClientOptions\n", + "from google.cloud import discoveryengine\n", + "\n", + "VERTEX_AI_SEARCH_LOCATION = \"global\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "kfkQMK1yBnN6" + }, + "source": [ + "### Create and Populate a Datastore" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": { + "id": "MwUPv8h5BqEJ" + }, + "outputs": [], + "source": [ + "def create_data_store(\n", + " project_id: str, location: str, data_store_name: str, data_store_id: str\n", + "):\n", + " # Create a client\n", + " client_options = (\n", + " ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n", + " if location != \"global\"\n", + " else None\n", + " )\n", + " client = discoveryengine.DataStoreServiceClient(client_options=client_options)\n", + "\n", + " # Initialize request argument(s)\n", + " data_store = discoveryengine.DataStore(\n", + " display_name=data_store_name,\n", + " industry_vertical=discoveryengine.IndustryVertical.GENERIC,\n", + " content_config=discoveryengine.DataStore.ContentConfig.CONTENT_REQUIRED,\n", + " )\n", + "\n", + " operation = client.create_data_store(\n", + " request=discoveryengine.CreateDataStoreRequest(\n", + " parent=client.collection_path(project_id, location, \"default_collection\"),\n", + " data_store=data_store,\n", + " data_store_id=data_store_id,\n", + " )\n", + " )\n", + "\n", + " # Make the request\n", + " response = operation.result(timeout=90)\n", + " return response.name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ABXeS6jCBs8H" + }, + "outputs": [], + "source": [ + "# The datastore name can only contain lowercase letters, numbers, and hyphens\n", + "DATASTORE_NAME = \"alphabet-contracts\" # @param {type:\"string\", isTemplate: true}\n", + "DATASTORE_ID = f\"{DATASTORE_NAME}-id\"\n", + "\n", + "create_data_store(PROJECT_ID, VERTEX_AI_SEARCH_LOCATION, DATASTORE_NAME, DATASTORE_ID)" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": { + "id": "jRMWKUJXBwhu" + }, + "outputs": [], + "source": [ + "def import_documents(\n", + " project_id: str,\n", + " location: str,\n", + " data_store_id: str,\n", + " gcs_uri: str,\n", + "):\n", + " # Create a client\n", + " client_options = (\n", + " ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n", + " if location != \"global\"\n", + " else None\n", + " )\n", + " client = discoveryengine.DocumentServiceClient(client_options=client_options)\n", + "\n", + " # The full resource name of the search engine branch.\n", + " # e.g. projects/{project}/locations/{location}/dataStores/{data_store_id}/branches/{branch}\n", + " parent = client.branch_path(\n", + " project=project_id,\n", + " location=location,\n", + " data_store=data_store_id,\n", + " branch=\"default_branch\",\n", + " )\n", + "\n", + " source_documents = [f\"{gcs_uri}/*\"]\n", + "\n", + " request = discoveryengine.ImportDocumentsRequest(\n", + " parent=parent,\n", + " gcs_source=discoveryengine.GcsSource(\n", + " input_uris=source_documents, data_schema=\"content\"\n", + " ),\n", + " # Options: `FULL`, `INCREMENTAL`\n", + " reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,\n", + " )\n", + "\n", + " # Make the request\n", + " operation = client.import_documents(request=request)\n", + "\n", + " response = operation.result()\n", + "\n", + " # Once the operation is complete,\n", + " # get information from operation metadata\n", + " metadata = discoveryengine.ImportDocumentsMetadata(operation.metadata)\n", + "\n", + " # Handle the response\n", + " return operation.operation.name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "fJlFkhe7BznS" + }, + "outputs": [], + "source": [ + "GCS_BUCKET = \"gs://cloud-samples-data/gen-app-builder/search/alphabet-investor-pdfs\" # @param {type:\"string\", isTemplate: true}\n", + "\n", + "import_documents(PROJECT_ID, VERTEX_AI_SEARCH_LOCATION, DATASTORE_ID, GCS_BUCKET)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "gvrbLtDj6zCv" + }, + "source": [ + "### Create a Search Engine" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": { + "id": "GHaRWVvZ6vOc" + }, + "outputs": [], + "source": [ + "def create_engine(\n", + " project_id: str, location: str, engine_name: str, engine_id: str, data_store_id: str\n", + "):\n", + " # Create a client\n", + " client_options = (\n", + " ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n", + " if location != \"global\"\n", + " else None\n", + " )\n", + " client = discoveryengine.EngineServiceClient(client_options=client_options)\n", + "\n", + " # Initialize request argument(s)\n", + " engine = discoveryengine.Engine(\n", + " display_name=engine_name,\n", + " solution_type=discoveryengine.SolutionType.SOLUTION_TYPE_SEARCH,\n", + " industry_vertical=discoveryengine.IndustryVertical.GENERIC,\n", + " data_store_ids=[data_store_id],\n", + " search_engine_config=discoveryengine.Engine.SearchEngineConfig(\n", + " search_tier=discoveryengine.SearchTier.SEARCH_TIER_ENTERPRISE,\n", + " ),\n", + " )\n", + "\n", + " request = discoveryengine.CreateEngineRequest(\n", + " parent=client.collection_path(project_id, location, \"default_collection\"),\n", + " engine=engine,\n", + " engine_id=engine.display_name,\n", + " )\n", + "\n", + " # Make the request\n", + " operation = client.create_engine(request=request)\n", + " response = operation.result(timeout=90)\n", + " return response.name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6Z21T3Pm8ngv" + }, + "outputs": [], + "source": [ + "ENGINE_NAME = DATASTORE_NAME\n", + "ENGINE_ID = DATASTORE_ID\n", + "create_engine(\n", + " PROJECT_ID, VERTEX_AI_SEARCH_LOCATION, ENGINE_NAME, ENGINE_ID, DATASTORE_ID\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "EdvJRUWRNGHE" + }, + "source": [ + "## Create a RAG corpus using Vertex AI Search Engine as the retrieval backend" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "5303c05f7aa6" + }, + "source": [ + "### Import libraries" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": { + "id": "6fc324893334" + }, + "outputs": [], + "source": [ + "from vertexai.preview import rag\n", + "from vertexai.preview.generative_models import GenerativeModel, Tool" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e43229f3ad4f" + }, + "source": [ + "### Create RAG Config\n" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": { + "id": "cf93d5f0ce00" + }, + "outputs": [], + "source": [ + "# Name your corpus\n", + "DISPLAY_NAME = \"\" # @param {type:\"string\", \"placeholder\": \"your-corpus-name\"}\n", + "\n", + "# Vertex AI Search name\n", + "ENGINE_NAME = \"\" # @param {type:\"string\", \"placeholder\": \"your-engine-name\"}\n", + "vertex_ai_search_config = rag.VertexAiSearchConfig(\n", + " serving_config=f\"{ENGINE_NAME}/servingConfigs/default_search\",\n", + ")\n", + "\n", + "rag_corpus = rag.create_corpus(\n", + " display_name=DISPLAY_NAME,\n", + " vertex_ai_search_config=vertex_ai_search_config,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2834a2721633" + }, + "outputs": [], + "source": [ + "# Check the corpus just created\n", + "new_corpus = rag.get_corpus(name=rag_corpus.name)\n", + "new_corpus" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e1c1c7944dc8" + }, + "source": [ + "## Using Gemini GenerateContent API with Rag Retrieval Tool\n" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": { + "id": "d3afdf13dbbb" + }, + "outputs": [], + "source": [ + "rag_resource = rag.RagResource(\n", + " rag_corpus=rag_corpus.name,\n", + ")\n", + "\n", + "rag_retrieval_tool = Tool.from_retrieval(\n", + " retrieval=rag.Retrieval(\n", + " source=rag.VertexRagStore(\n", + " rag_resources=[rag_resource], # Currently only 1 corpus is allowed.\n", + " similarity_top_k=10,\n", + " ),\n", + " )\n", + ")\n", + "\n", + "rag_model = GenerativeModel(\"gemini-1.5-flash\", tools=[rag_retrieval_tool])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "hLAzPNJOOKuI" + }, + "source": [ + "Note: The Vertex AI Search engine will take some time to be ready to query.\n", + "\n", + "If you recently created an engine and you receive an error similar to:\n", + "\n", + "`404 Engine {ENGINE_NAME} is not found`\n", + "\n", + "Then wait a few minutes and try your query again." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "484f5242dcae" + }, + "outputs": [], + "source": [ + "GENERATE_CONTENT_PROMPT = (\n", + " \"Who is CFO of Google?\" # @param {type:\"string\", isTemplate: true}\n", + ")\n", + "\n", + "response = rag_model.generate_content(GENERATE_CONTENT_PROMPT)\n", + "\n", + "response" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "287a90fed14f" + }, + "source": [ + "## Using other generation API with Rag Retrieval Tool\n", + "\n", + "The retrieved contexts can be passed to any SDK or model generation API to generate final results." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "428921dea97d" + }, + "outputs": [], + "source": [ + "RETRIEVAL_QUERY = \"Who is CFO of Google?\" # @param {type:\"string\", isTemplate: true}\n", + "\n", + "rag_resource = rag.RagResource(rag_corpus=rag_corpus.name)\n", + "\n", + "response = rag.retrieval_query(\n", + " rag_resources=[rag_resource], # Currently only 1 corpus is allowed.\n", + " text=RETRIEVAL_QUERY,\n", + " similarity_top_k=10,\n", + ")\n", + "\n", + "# The retrieved context can be passed to any SDK or model generation API to generate final results.\n", + "retrieved_context = \" \".join(\n", + " [context.text for context in response.contexts.contexts]\n", + ").replace(\"\\n\", \"\")\n", + "\n", + "retrieved_context" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "2a4e033321ad" + }, + "source": [ + "## Cleaning up\n", + "\n", + "Clean up RAG resources created in this notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "a105caffd9e7" + }, + "outputs": [], + "source": [ + "delete_rag_corpus = False # @param {type:\"boolean\"}\n", + "\n", + "if delete_rag_corpus:\n", + " rag.delete_corpus(name=rag_corpus.name)" + ] + } + ], + "metadata": { + "colab": { + "name": "rag_engine_vertex_ai_search.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/rag-grounding/README.md b/rag-grounding/README.md index 3bdcf92ba5..5e4eb86f0c 100644 --- a/rag-grounding/README.md +++ b/rag-grounding/README.md @@ -100,6 +100,7 @@ data. - **[rag_engine_pinecone.ipynb](../gemini/rag-engine/rag_engine_pinecone.ipynb)**: How to use RAG Engine with [Pinecone](https://www.pinecone.io/) - **[rag_engine_vector_search.ipynb](../gemini/rag-engine/rag_engine_vector_search.ipynb)**: How to use RAG Engine with [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview) - **[rag_engine_weaviate.ipynb](../gemini/rag-engine/rag_engine_weaviate.ipynb)**: How to use RAG Engine with [Weaviate](https://weaviate.io/) +- **[rag_engine_vertex_ai_search.ipynb](../gemini/rag-engine/rag_engine_vertex_ai_search.ipynb)**: How to use RAG Engine with [Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/enterprise-search-introduction) - **[rag_engine_evaluation.ipynb](../gemini/rag-engine/rag_engine_evaluation.ipynb)**: Advanced RAG Techniques and Evaluation with RAG Engine ## Bring your own Search for RAG/Grounding