mirrors/googlecloudplatform--generative-ai
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:30:30 +08:00
commit 914fea506e
2793 changed files with 802106 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gemini/rag-engine/intro_rag_engine.ipynb
+580
View File
@@ -0,0 +1,580 @@
{
"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": [
"# Intro to Building a Scalable and Modular RAG System with RAG Engine on Gemini Enterprise Agent Platform\n",
"\n",
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/agent-platform/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Fintro_rag_engine.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/agent-platform/workbench/instances?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/intro_rag_engine.ipynb\">\n",
" <img width=\"32px\" src=\"https://storage.googleapis.com/github-repo/workbench-icon.svg\" alt=\"Workbench logo\"><br> Open in Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<p>\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a>\n",
"</p> "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "84f0f73a0f76"
},
"source": [
"| Author |\n",
"| --- |\n",
"| [Holt Skinner](https://github.com/holtskinner) |"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tvgnzT1CKxrO"
},
"source": [
"## Overview\n",
"\n",
"Retrieval Augmented Generation (RAG) improves Large Language Models (LLMs) by allowing them to access and process external information sources during generation. This ensures the model's responses are grounded in factual data and avoids hallucinations.\n",
"\n",
"A common problem with LLMs is that they don't understand private knowledge, that\n",
"is, your organization's data. With RAG Engine, you can enrich the\n",
"LLM context with additional private information, because the model can reduce\n",
"hallucinations and answer questions more accurately.\n",
"\n",
"By combining additional knowledge sources with the existing knowledge that LLMs\n",
"have, a better context is provided. The improved context along with the query\n",
"enhances the quality of the LLM's response.\n",
"\n",
"The following concepts are key to understanding Agent Platform RAG Engine. These concepts are listed in the order of the\n",
"retrieval-augmented generation (RAG) process.\n",
"\n",
"1. **Data ingestion**: Intake data from different data sources. For example,\n",
" local files, Google Cloud Storage, and Google Drive.\n",
"\n",
"1. **Data transformation**: Conversion of the data in preparation for indexing. For example, data is split into chunks.\n",
"\n",
"1. **Embedding**: Numerical representations of words or pieces of text. These numbers capture the\n",
" semantic meaning and context of the text. Similar or related words or text\n",
" tend to have similar embeddings, which means they are closer together in the\n",
" high-dimensional vector space.\n",
"\n",
"1. **Data indexing**: RAG Engine creates an index called a corpus.\n",
" The index structures the knowledge base so it's optimized for searching. For\n",
" example, the index is like a detailed table of contents for a massive\n",
" reference book.\n",
"\n",
"1. **Retrieval**: When a user asks a question or provides a prompt, the retrieval\n",
" component in RAG Engine searches through its knowledge\n",
" base to find information that is relevant to the query.\n",
"\n",
"1. **Generation**: The retrieved information becomes the context added to the\n",
" original user query as a guide for the generative AI model to generate\n",
" factually grounded and relevant responses.\n",
"\n",
"For more information, refer to the public documentation for [Agent Platform RAG Engine](https://docs.cloud.google.com/gemini-enterprise-agent-platform/build/rag-engine/rag-overview)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "61RBz8LLbxCR"
},
"source": [
"## Get started"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "No17Cw5hgx12"
},
"source": [
"### Install Agent Platform SDK and Google Gen AI SDK\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tFy3H3aPgx12"
},
"outputs": [],
"source": [
"%pip install --upgrade --quiet google-cloud-aiplatform google-genai"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dmWOrTJ3gx13"
},
"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": null,
"metadata": {
"id": "NyKGtVQjgx13"
},
"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": "DF4l8DTdWgPY"
},
"source": [
"### Set Google Cloud project information\n",
"\n",
"To get started using Agent Platform, you must have an existing Google Cloud project and [enable the Agent Platform API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
"\n",
"Learn more about [setting up a project](https://docs.cloud.google.com/resource-manager/docs/creating-managing-projects) and a [development environment](https://cloud.google.com/docs/authentication/set-up-adc-local-dev-environment).\n",
"\n",
"See [supported regions](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-engine/rag-overview#supported-regions) for location options."
]
},
{
"cell_type": "code",
"execution_count": null,
"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 agentplatform\n",
"\n",
"# fmt: off\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
"# fmt: on\n",
"if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
" PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
"\n",
"# See https://cloud.google.com/vertex-ai/generative-ai/docs/rag-engine/rag-overview#supported-regions for location options.\n",
"client = agentplatform.Client(project=PROJECT_ID, location=\"us-east1\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5303c05f7aa6"
},
"source": [
"### Import libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6fc324893334"
},
"outputs": [],
"source": [
"from IPython.display import Markdown, display\n",
"from agentplatform import types\n",
"from google import genai\n",
"from google.genai import types as genai_types"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "e43229f3ad4f"
},
"source": [
"### Create a RAG Corpus"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cf93d5f0ce00"
},
"outputs": [],
"source": [
"rag_corpus = client.rag.create_corpus(\n",
" rag_corpus=types.RagCorpus(\n",
" display_name=\"my-rag-corpus\",\n",
" )\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "197c585b61b2"
},
"source": [
"### Check the corpus just created"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "f229b13dc617"
},
"outputs": [],
"source": [
"client.rag.list_corpora()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "c52924cc1440"
},
"source": [
"### Upload a local file to the corpus"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4976ffe8564f"
},
"outputs": [],
"source": [
"%%writefile test.md\n",
"\n",
"Retrieval-Augmented Generation (RAG) is a technique that enhances the capabilities of large language models (LLMs) by allowing them to access and incorporate external data sources when generating responses. Here's a breakdown:\n",
"\n",
"**What it is:**\n",
"\n",
"* **Combining Retrieval and Generation:**\n",
" * RAG combines the strengths of information retrieval systems (like search engines) with the generative power of LLMs.\n",
" * It enables LLMs to go beyond their pre-trained data and access up-to-date and specific information.\n",
"* **How it works:**\n",
" * When a user asks a question, the RAG system first retrieves relevant information from external data sources (e.g., databases, documents, web pages).\n",
" * This retrieved information is then provided to the LLM as additional context.\n",
" * The LLM uses this augmented context to generate a more accurate and informative response.\n",
"\n",
"**Why it's helpful:**\n",
"\n",
"* **Access to Up-to-Date Information:**\n",
" * LLMs are trained on static datasets, so their knowledge can become outdated. RAG allows them to access real-time or frequently updated information.\n",
"* **Improved Accuracy and Factual Grounding:**\n",
" * RAG reduces the risk of LLM \"hallucinations\" (generating false or misleading information) by grounding responses in verified external data.\n",
"* **Enhanced Contextual Relevance:**\n",
" * By providing relevant context, RAG enables LLMs to generate more precise and tailored responses to specific queries.\n",
"* **Increased Trust and Transparency:**\n",
" * RAG can provide source citations, allowing users to verify the information and increasing trust in the LLM's responses.\n",
"* **Cost Efficiency:**\n",
" * Rather than constantly retraining large language models, RAG allows for the introduction of new data in a more cost effective way.\n",
"\n",
"In essence, RAG bridges the gap between the vast knowledge of LLMs and the need for accurate, current, and contextually relevant information.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "529390917c29"
},
"outputs": [],
"source": [
"rag_file = client.rag.upload_file(\n",
" corpus_name=rag_corpus.name,\n",
" path=\"test.md\",\n",
" display_name=\"test.md\",\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5269a0c2786d"
},
"source": [
"### Import files from Google Cloud Storage\n",
"\n",
"Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Google Cloud Storage bucket.\n",
"\n",
"For this example, we'll use a public GCS bucket containing earning reports from Alphabet."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5910ae450f69"
},
"outputs": [],
"source": [
"INPUT_GCS_BUCKET = (\n",
" \"gs://cloud-samples-data/gen-app-builder/search/alphabet-investor-pdfs/*\"\n",
")\n",
"\n",
"response = client.rag.import_files(\n",
" name=rag_corpus.name,\n",
" import_config=types.ImportRagFilesConfig(\n",
" gcs_source=types.GcsSource(uris=[INPUT_GCS_BUCKET]),\n",
" rag_file_transformation_config=types.RagFileTransformationConfig(\n",
" rag_file_chunking_config=types.RagFileChunkingConfig(\n",
" chunk_size=1024,\n",
" chunk_overlap=100,\n",
" )\n",
" ), # optional\n",
" max_embedding_requests_per_min=900, # optional\n",
" ),\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "60a84095746d"
},
"source": [
"### Import files from Google Drive\n",
"\n",
"Eligible paths can be formatted as:\n",
"\n",
"- `https://drive.google.com/drive/folders/{folder_id}`\n",
"- `https://drive.google.com/file/d/{file_id}`.\n",
"\n",
"Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Drive folder/files.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0a90c125874c"
},
"outputs": [],
"source": [
"response = client.rag.import_files(\n",
" name=rag_corpus.name,\n",
" import_config=types.ImportRagFilesConfig(\n",
" google_drive_source=types.GoogleDriveSource(\n",
" resource_ids=[\n",
" types.GoogleDriveSourceResourceId(\n",
" resource_id=\"https://drive.google.com/drive/folders/{folder_id}\",\n",
" resource_type=types.ResourceType.RESOURCE_TYPE_FOLDER,\n",
" )\n",
" ]\n",
" ),\n",
" rag_file_transformation_config=types.RagFileTransformationConfig(\n",
" rag_file_chunking_config=types.RagFileChunkingConfig(\n",
" chunk_size=512,\n",
" chunk_overlap=50,\n",
" )\n",
" ), # optional\n",
" ),\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "f700b3e23121"
},
"source": [
"### Optional: Perform direct context retrieval"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4669c5cdbb5a"
},
"outputs": [],
"source": [
"# Direct context retrieval\n",
"response = client.rag.retrieve_contexts(\n",
" vertex_rag_store=genai_types.VertexRagStore(\n",
" rag_resources=[\n",
" genai_types.VertexRagStoreRagResource(rag_corpus=rag_corpus.name)\n",
" ]\n",
" ),\n",
" query=types.RagQuery(\n",
" text=\"What is RAG and why it is helpful?\",\n",
" rag_retrieval_config=types.RagRetrievalConfig(\n",
" top_k=10,\n",
" filter=types.RagRetrievalConfigFilter(vector_distance_threshold=0.5),\n",
" ),\n",
" ),\n",
")\n",
"print(response)\n",
"\n",
"# Optional: The retrieved context can be passed to any SDK or model generation API to generate final results.\n",
"# context = \" \".join([context.text for context in response.contexts.contexts]).replace(\"\\n\", \"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "79ea89661842"
},
"source": [
"### Create RAG Retrieval Tool"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0ebceac3d816"
},
"outputs": [],
"source": [
"# Create a tool for the RAG Corpus\n",
"rag_retrieval_tool = genai_types.Tool(\n",
" retrieval=genai_types.Retrieval(\n",
" vertex_rag_store=genai_types.VertexRagStore(\n",
" rag_corpora=[rag_corpus.name],\n",
" similarity_top_k=10,\n",
" vector_distance_threshold=0.5,\n",
" )\n",
" )\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "d88fa7ede853"
},
"source": [
"### Generate Content with Gemini using RAG Retrieval Tool"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8dd928baecd4"
},
"outputs": [],
"source": [
"MODEL_ID = \"gemini-3.5-flash\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1f02122723b6"
},
"outputs": [],
"source": [
"# Create a Gen AI SDK client\n",
"genai_client = genai.Client(project=PROJECT_ID, location=\"your-location\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "124b36be8d5b"
},
"outputs": [],
"source": [
"response = genai_client.models.generate_content(\n",
" model=MODEL_ID,\n",
" contents=\"What is RAG?\",\n",
" config=genai_types.GenerateContentConfig(tools=[rag_retrieval_tool]),\n",
")\n",
"\n",
"display(Markdown(response.text))"
]
}
],
"metadata": {
"colab": {
"name": "intro_rag_engine.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
gemini/rag-engine/rag_engine_cross_corpus_retrieval.ipynb
+368
View File
@@ -0,0 +1,368 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "T563aA8BZ-Nk"
},
"outputs": [],
"source": [
"# Copyright 2026 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": "TE2sLvrjmNx7"
},
"source": [
"# Vertex AI Rag: Cross-Corpus Retrieval with `AskContexts` and `AsyncRetrieveContexts` Demo\n",
"\n",
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/cross_corpus_retrieval.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Fcross_corpus_retrieval.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/cross_corpus_retrieval.ipynb\">\n",
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/cross_corpus_retrieval.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/cross_corpus_retrieval.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/cross_corpus_retrieval.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/cross_corpus_retrieval.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/cross_corpus_retrieval.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/cross_corpus_retrieval.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a> "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"| Author |\n",
"| --- |\n",
"| [Tiger Jin](https://github.com/tiger-w-jin) |"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "f9de46eb"
},
"source": [
"\n",
"\n",
"This notebook demonstrates how to invoke the new Vertex AI RAG handlers (`AskContexts` and `AsyncRetrieveContexts`) using the Vertex AI Python SDK for cross-corpus retrieval.\n",
"\n",
"**Note**: Since these handlers are currently in preview states, this notebook relies on the `vertexai.preview.rag` module."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "b0f9ed43"
},
"source": [
"## Cross-Corpus Retrieval\n",
"\n",
"Cross-corpus retrieval in Vertex AI RAG allows you to query multiple RAG corpora simultaneously with a single request. This is particularly powerful when your knowledge base is distributed across different datasets or when you need to retrieve information that might span across various domains or document types.\n",
"\n",
"For more information, refer to the [official documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-engine/cross-corpus-retrieval)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "1f78c0b3"
},
"source": [
"### `AskContexts` and `AsyncRetrieveContexts`\n",
"\n",
"These are two key functions within the Vertex AI RAG framework, available through the `vertexai.preview.rag` module. They are designed to facilitate the retrieval of relevant contexts from configured RAG corpora.\n",
"\n",
"* **`AskContexts`**: This function is used for synchronous retrieval of contexts. It takes a query and a list of `RagResource` objects (which define the RAG corpora to search) and returns the most relevant contexts based on the `rag_retrieval_config`.\n",
"\n",
"* **`AsyncRetrieveContexts`**: Similar to `AskContexts`, but designed for asynchronous retrieval. This is particularly useful when dealing with multiple corpora or when integration into an asynchronous application flow is required, allowing for non-blocking operations."
]
},
{
"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": "f771a76b"
},
"outputs": [],
"source": [
"%pip install google-cloud-aiplatform"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dmWOrTJ3gx13"
},
"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": null,
"metadata": {
"id": "5cd723c4"
},
"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": "new_init_heading"
},
"source": [
"## Initialization\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": null,
"metadata": {
"id": "390b571d"
},
"outputs": [],
"source": [
"from google.cloud import aiplatform\n",
"import vertexai\n",
"from vertexai.preview import rag\n",
"from google.cloud import aiplatform_v1beta1\n",
"\n",
"SERVICE = \"aiplatform.googleapis.com\" # @param {type:\"string\"}\n",
"ENDPOINT = \"us-central1-aiplatform.googleapis.com\" # @param {type:\"string\"}\n",
"\n",
"\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n",
"LOCATION = \"us-central1\" # @param {type:\"string\"}\n",
"\n",
"aiplatform.init(project=PROJECT_ID, location=LOCATION,\n",
" api_endpoint=ENDPOINT\n",
")\n",
"print(f\"Vertex AI initialized with project: {PROJECT_ID}, location: {LOCATION}\")\n",
"print(\"Available RAG functions exposed:\")\n",
"print(f\" - async_retrieve_contexts: {hasattr(rag, 'async_retrieve_contexts')}\")\n",
"print(f\" - ask_contexts: {hasattr(rag, 'ask_contexts')}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "new_corpus_heading"
},
"source": [
"## Rag Corpus management\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Option 1: Create a RAG Corpus\n",
"\n",
"To create a new RAG corpus, use the `create_corpus` function from the `rag` module. This initializes a managed index for your documents. You can reference the [introductory notebook](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb) for more details."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "d4b1a027"
},
"source": [
"### Option 2: Use Existing Rag Corpus Path\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "74acf5ea"
},
"outputs": [],
"source": [
"CORPUS_PATH_1 = \"[your-corpus-path]\" # @param {type:\"string\"}\n",
"CORPUS_PATH_2 = \"[your-corpus-path]\" # @param {type:\"string\"}\n",
"print(f\"Corpus 1 set to: {CORPUS_PATH_1}\")\n",
"print(f\"Corpus 2 set to: {CORPUS_PATH_2}\")\n",
"\n",
"rag_resources_list = [\n",
" rag.RagResource(rag_corpus=CORPUS_PATH_1),\n",
" rag.RagResource(rag_corpus=CORPUS_PATH_2),\n",
"]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "new_retrieval_heading"
},
"source": [
"## Rag Retrieval\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "077a2d20"
},
"source": [
"### AskContexts\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "M0YHP3BhG6tP"
},
"outputs": [],
"source": [
"# # Config\n",
"config = rag.RagRetrievalConfig(top_k=2)\n",
"\n",
"query=\"[Your query]\"# @param {type:\"string\"}\n",
"\n",
"response = rag.ask_contexts(\n",
" text=query,\n",
" rag_resources=rag_resources_list,\n",
" rag_retrieval_config=config,\n",
")\n",
"\n",
"print(\"Response:\", response)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "89bb9987"
},
"source": [
"### AsyncRetrieveContexts\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0f8098d7"
},
"outputs": [],
"source": [
"query=\"[Your query]\"# @param {type:\"string\"}\n",
"\n",
"# Async retrieve contexts from multiple corpora\n",
"response = await rag.async_retrieve_contexts(\n",
" text=query,\n",
" rag_resources=rag_resources_list,\n",
" rag_retrieval_config=config,\n",
")\n",
"print(response)"
]
}
],
"metadata": {
"colab": {
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
gemini/rag-engine/rag_engine_eval_service_sdk.ipynb
+738
View File
@@ -0,0 +1,738 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ur8xi4C7S06n"
},
"outputs": [],
"source": [
"# Copyright 2025 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": [
"# Evaluating Vertex RAG Engine Generation with Vertex AI Python SDK for Gen AI Evaluation Service\n",
"\n",
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Frag_engine_eval_service_sdk.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb\">\n",
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb\">\n",
" <img width=\"32px\" src=\"https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/53/X_logo_2023_original.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_eval_service_sdk.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "84f0f73a0f76"
},
"source": [
"| | |\n",
"|-|-|\n",
"| Author(s) | [Noa Ben-Efraim](https://github.com/noabenefraim/) |"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tvgnzT1CKxrO"
},
"source": [
"## Overview\n",
"\n",
"This notebook demonstrates how to evaluate the performance of a Retrieval Augmented Generation (RAG) engine built with Vertex AI using the Vertex AI Python SDK for Gen AI Evaluation Service. By focusing on a practical example using \"Alice in Wonderland\" as our knowledge base, we'll walk through the process of creating an evaluation dataset and applying custom metrics to assess the quality of generated responses.\n",
"\n",
"Specifically, this notebook will guide you through:\n",
"\n",
"* **Setting up a RAG Corpus:** Creating and populating a RAG corpus with a PDF document.\n",
"* **Generating Grounded Responses:** Using the Vertex AI Gemini model to produce responses based on retrieved contexts.\n",
"* **Creating an Evaluation Dataset:** Constructing a dataset with prompts, retrieved contexts, and generated responses.\n",
"* **Defining Custom Evaluation Metrics:** Implementing a custom metric to assess the accuracy, completeness, and groundedness of the generated responses.\n",
"* **Running Evaluation Tasks:** Utilizing the Vertex AI Gen AI Evaluation Service to evaluate the RAG engine's performance.\n",
"* **Analyzing Evaluation Results:** Visualizing and interpreting the evaluation results using the provided SDK tools."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "61RBz8LLbxCR"
},
"source": [
"## Get started"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "No17Cw5hgx12"
},
"source": [
"### Install Google Gen AI SDK and other required packages\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tFy3H3aPgx12"
},
"outputs": [],
"source": [
"%pip install --upgrade --quiet google-genai google-cloud-aiplatform[evaluation] vertexai"
]
},
{
"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": [
"<div class=\"alert alert-block alert-warning\">\n",
"<b>⚠️ The kernel is going to restart. In Colab or Colab Enterprise, you might see an error message that says \"Your session crashed for an unknown reason.\" This is expected. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
"</div>\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dmWOrTJ3gx13"
},
"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": null,
"metadata": {
"id": "NyKGtVQjgx13"
},
"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": "DF4l8DTdWgPY"
},
"source": [
"### Set Google Cloud project information\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": "Nqwi-5ufWp_B"
},
"outputs": [],
"source": [
"# Use the environment variable if the user doesn't provide Project ID.\n",
"import os\n",
"\n",
"from google import genai\n",
"import vertexai\n",
"\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
"if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
" PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
"LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
"\n",
"client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)\n",
"vertexai.init(project=PROJECT_ID, location=LOCATION)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5303c05f7aa6"
},
"source": [
"### Import libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6fc324893334"
},
"outputs": [],
"source": [
"from google.genai.types import GenerateContentConfig, Retrieval, Tool, VertexRagStore\n",
"import pandas as pd\n",
"from tqdm import tqdm\n",
"from vertexai import rag\n",
"from vertexai.evaluation import (\n",
" EvalTask,\n",
" PointwiseMetric,\n",
" PointwiseMetricPromptTemplate,\n",
" notebook_utils,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "e43229f3ad4f"
},
"source": [
"### Load model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cf93d5f0ce00"
},
"outputs": [],
"source": [
"MODEL_ID = \"gemini-2.5-flash\" # @param {type:\"string\"}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "a45112933981"
},
"source": [
"### Create `RAGCorpus`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1b1ba2900db9"
},
"outputs": [],
"source": [
"# Currently supports Google first-party embedding models\n",
"EMBEDDING_MODEL = \"publishers/google/models/text-embedding-005\" # @param {type:\"string\", isTemplate: true}\n",
"\n",
"rag_corpus = rag.create_corpus(\n",
" display_name=\"rag-eval-corpus\",\n",
" description=\"A test corpus for generation evaluation\",\n",
" backend_config=rag.RagVectorDbConfig(\n",
" rag_embedding_model_config=rag.RagEmbeddingModelConfig(\n",
" vertex_prediction_endpoint=rag.VertexPredictionEndpoint(\n",
" publisher_model=EMBEDDING_MODEL\n",
" )\n",
" )\n",
" ),\n",
")\n",
"\n",
"# Get the rag corpus you just created\n",
"rag.get_corpus(rag_corpus.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5a1cbd3d0f96"
},
"source": [
"### Import files from Google Cloud Storage into `RagCorpus` (configure chunk size, chunk overlap etc as desired)\n",
"\n",
"For this step you will need to create a GCS bucket, and then copy over the data from the public GCS bucket. Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com) for your Google Cloud Storage bucket.\n",
"\n",
"For this example, we'll use a dataset that comprises the full texts of five classic children's literature books: \"The Wizard of Oz,\" \"Gulliver's Travels,\" \"Peter Pan,\" \"Alice's Adventures in Wonderland,\" and \"Through the Looking-Glass.\" This collection provides a rich corpus for exploring themes, characters, and settings across these iconic stories.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "20d44cdad6ab"
},
"source": [
"##### Copy data from public GCS bucket"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ca8381cc8a9b"
},
"outputs": [],
"source": [
"CURRENT_BUCKET_PATH = \"gs://\" # @param {type:\"string\"},\n",
"\n",
"PUBLIC_DATA_PATH = (\n",
" \"gs://github-repo/generative-ai/gemini/rag-engine/rag_engine_eval_service/\"\n",
")\n",
"\n",
"!gsutil -m rsync -r -d $PUBLIC_DATA_PATH $CURRENT_BUCKET_PATH"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "b84020847859"
},
"source": [
"##### Import dataset into `RagCorpus`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8f10ecc2c035"
},
"outputs": [],
"source": [
"transformation_config = rag.TransformationConfig(\n",
" chunking_config=rag.ChunkingConfig(\n",
" chunk_size=512,\n",
" chunk_overlap=100,\n",
" ),\n",
")\n",
"\n",
"rag.import_files(\n",
" corpus_name=rag_corpus.name,\n",
" paths=[CURRENT_BUCKET_PATH],\n",
" transformation_config=transformation_config, # Optional\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "609cdfd973c6"
},
"outputs": [],
"source": [
"# List the files in the rag corpus\n",
"rag.list_files(rag_corpus.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dbb9974409ca"
},
"source": [
"### Create RAG Retrieval Tool"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "d767af7c494e"
},
"outputs": [],
"source": [
"# Create a tool for the RAG Corpus\n",
"rag_retrieval_tool = Tool(\n",
" retrieval=Retrieval(\n",
" vertex_rag_store=VertexRagStore(\n",
" rag_corpora=[rag_corpus.name],\n",
" similarity_top_k=10,\n",
" vector_distance_threshold=0.5,\n",
" )\n",
" )\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ce3f57ffd049"
},
"outputs": [],
"source": [
"def get_generated_response(prompt: str) -> str:\n",
" \"\"\"\n",
" Generates a grounded response using a language model and retrieved context.\n",
"\n",
" Args:\n",
" prompt: The input prompt for the language model.\n",
"\n",
" Returns:\n",
" The generated text response.\n",
" \"\"\"\n",
" response = client.models.generate_content(\n",
" model=MODEL_ID,\n",
" contents=prompt,\n",
" config=GenerateContentConfig(tools=[rag_retrieval_tool]),\n",
" )\n",
"\n",
" return response.text"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "bc928efd7e38"
},
"outputs": [],
"source": [
"def get_retrieved_contexts(prompt: str) -> str:\n",
" \"\"\"\n",
" Retrieves relevant contexts based on a given prompt using a RAG system.\n",
"\n",
" Args:\n",
" prompt: The input prompt for context retrieval.\n",
"\n",
" Returns:\n",
" A concatenated string of retrieved context texts, with newlines removed.\n",
" \"\"\"\n",
"\n",
" rag_filter = rag.utils.resources.Filter(vector_distance_threshold=0.5)\n",
"\n",
" retrieval_config = rag.RagRetrievalConfig(top_k=5, filter=rag_filter)\n",
"\n",
" response = rag.retrieval_query(\n",
" rag_resources=[\n",
" rag.RagResource(\n",
" rag_corpus=rag_corpus.name,\n",
" # Optional: supply IDs from `rag.list_files()`.\n",
" # rag_file_ids=[\"rag-file-1\", \"rag-file-2\", ...],\n",
" )\n",
" ],\n",
" text=prompt,\n",
" rag_retrieval_config=retrieval_config,\n",
" )\n",
" context = \" \".join(\n",
" [context.text for context in response.contexts.contexts]\n",
" ).replace(\"\\n\", \"\")\n",
" return context"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8a904150ee7d"
},
"source": [
"### Create Evaluation Dataset\n",
"\n",
"Now we are prepared to create the evaluation dataset. The dataset will include:\n",
"\n",
"+ Prompt: What the user is asking the RAG engine. The prompts will be a mix of inter-document and intra-document analysis.\n",
"+ Retrieved Context: The top k retrieved context from Vertex RAG Engine\n",
"+ Generated Response: The LLM generated responses grounded in the retrieved context."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "00ffc51df8b6"
},
"outputs": [],
"source": [
"prompts = [\n",
" \"Compare and contrast the behaviors of the Mad Hatter and the March Hare during the tea party.\",\n",
" \"What happened during Alice's croquet game with the Queen of Hearts?\",\n",
" \"How did the Mad Hatter and March Hare act at the tea party?\",\n",
" \"What was special about the cakes Alice ate?\",\n",
" \"What happened when Gulliver first arrived in Lilliput?\",\n",
" \"What was Captain Hook's main goal in Neverland?\",\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "531d827a83fd"
},
"outputs": [],
"source": [
"retrieved_context = []\n",
"generated_response = []\n",
"for prompt in tqdm(prompts):\n",
" retrieved_context.append(get_retrieved_contexts(prompt))\n",
" generated_response.append(get_generated_response(prompt))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0d2d7a3f0b76"
},
"outputs": [],
"source": [
"eval_dataset = pd.DataFrame(\n",
" {\n",
" \"prompt\": prompts,\n",
" \"retrieved_context\": retrieved_context,\n",
" \"response\": generated_response,\n",
" }\n",
")\n",
"\n",
"eval_dataset"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ad1d81adc1cd"
},
"source": [
"## Use Gen AI Evaluation Service SDK\n",
"\n",
"Before diving into the evaluation process, we've set up the necessary components: a RAG corpus containing our document, a retrieval tool, and functions to generate grounded responses and retrieve relevant contexts. We've also compiled an evaluation dataset with a set of questions, the corresponding retrieved contexts, and the model's responses.\n",
"\n",
"This dataset will serve as the foundation for our evaluation. We'll now leverage the Vertex AI Gen AI Evaluation Service SDK to define and apply custom metrics, allowing us to quantitatively assess the RAG engine's performance. The Gen AI Evaluation Service provides a robust framework for creating and running evaluation tasks, enabling us to gain valuable insights into the quality of our generated responses."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "c056fb58f6e1"
},
"outputs": [],
"source": [
"custom_question_answering_correctness = PointwiseMetric(\n",
" metric=\"custom_question_answering_correctness\",\n",
" metric_prompt_template=PointwiseMetricPromptTemplate(\n",
" criteria={\n",
" \"accuracy\": (\n",
" \"The response provides completely accurate information, consistent with the retrieved context, with no errors or omissions.\"\n",
" ),\n",
" \"completeness\": (\n",
" \"The response answers all parts of the question fully, utilizing the information available in the retrieved context.\"\n",
" ),\n",
" \"groundedness\": (\n",
" \"The response uses only the information provided in the retrieved context and does not introduce any external information or hallucinations.\"\n",
" ),\n",
" },\n",
" rating_rubric={\n",
" \"5\": \"(Very good). The answer is completely accurate, complete, concise, grounded in the retrieved context, and follows all instructions.\",\n",
" \"4\": \"(Good). The answer is mostly accurate, complete, and grounded in the retrieved context, with minor issues in conciseness or instruction following.\",\n",
" \"3\": \"(Ok). The answer is partially accurate and complete but may have some inaccuracies, omissions, or significant issues with conciseness, groundedness, or instruction following, based on the retrieved context.\",\n",
" \"2\": \"(Bad). The answer contains significant inaccuracies, is largely incomplete, or fails to follow key instructions, considering the information available in the retrieved context.\",\n",
" \"1\": \"(Very bad). The answer is completely inaccurate, irrelevant, or fails to address the question in any meaningful way, based on the retrieved context.\",\n",
" },\n",
" input_variables=[\"prompt\", \"retrieved_context\"],\n",
" ),\n",
")\n",
"\n",
"# Display the serialized metric prompt template\n",
"print(custom_question_answering_correctness.metric_prompt_template)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bc6cd023895a"
},
"source": [
"### Run Eval Task\n",
"\n",
"The Gen AI Evaluation SDK has many useful utilities to graph, summarize, and explain the evaluation results. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "bbbe1f16e57e"
},
"outputs": [],
"source": [
"# Run evaluation using the custom_text_quality metric\n",
"eval_task = EvalTask(\n",
" dataset=eval_dataset,\n",
" metrics=[custom_question_answering_correctness],\n",
" experiment=\"test\",\n",
")\n",
"eval_result = eval_task.evaluate()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0e23de09f2b8"
},
"outputs": [],
"source": [
"notebook_utils.display_eval_result(eval_result=eval_result)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2fb80150675a"
},
"outputs": [],
"source": [
"# Example for graphing\n",
"notebook_utils.display_radar_plot(\n",
" eval_results_with_title=[(\"Question answering correctness\", eval_result)],\n",
" metrics=[\"custom_question_answering_correctness\"],\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9a40053a479e"
},
"outputs": [],
"source": [
"# Displaying explanations for one row.\n",
"notebook_utils.display_explanations(eval_result=eval_result, num=1)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2a4e033321ad"
},
"source": [
"## Cleaning up\n",
"\n",
"Delete ExperimentRun created by the evaluation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "113ab9d8887a"
},
"outputs": [],
"source": [
"aiplatform.ExperimentRun(\n",
" run_name=eval_result.metadata[\"experiment_run\"],\n",
" experiment=eval_result.metadata[\"experiment\"],\n",
").delete()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "db1bfe14c505"
},
"outputs": [],
"source": [
"rag.delete_corpus(rag_corpus.name)"
]
}
],
"metadata": {
"colab": {
"name": "rag_engine_eval_service_sdk.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
gemini/rag-engine/rag_engine_evaluation.ipynb
+1257
View File
File diff suppressed because it is too large Load Diff
gemini/rag-engine/rag_engine_feature_store.ipynb
+716
View File
@@ -0,0 +1,716 @@
{
"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 Feature Store\n",
"\n",
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_feature_store.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Frag_engine_feature_store.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/rag_engine_feature_store.ipynb\">\n",
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_feature_store.ipynb\">\n",
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_feature_store.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_feature_store.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_feature_store.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_feature_store.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_feature_store.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_feature_store.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a> "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "84f0f73a0f76"
},
"source": [
"| Author |\n",
"| --- |\n",
"| [Holt Skinner](https://github.com/holtskinner) |"
]
},
{
"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 Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/latest/overview) as a vector database.\n",
"\n",
"RAG Engine uses a built-in vector database powered by Spanner to store and manage vector representations of text documents. The vector database retrieves relevant documents based on the documents' semantic similarity to a given query.\n",
"\n",
"By integrating Vertex AI Feature Store as an additional vector database, RAG Engine can use Vertex AI Feature Store to handle large data volumes with low latency, which helps to improve the performance and scalability of your RAG applications.\n",
"\n",
"For more information, refer to the [official documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/use-vertexai-vector-search).\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-genai google-cloud-bigquery"
]
},
{
"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": [
"<div class=\"alert alert-block alert-warning\">\n",
"<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
"</div>\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dmWOrTJ3gx13"
},
"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": null,
"metadata": {
"id": "NyKGtVQjgx13"
},
"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": "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": null,
"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",
"from google import genai\n",
"\n",
"# fmt: off\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
"# fmt: on\n",
"if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
" PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
"\n",
"# See https://cloud.google.com/vertex-ai/generative-ai/docs/rag-engine/rag-overview#supported-regions for location options.\n",
"vertexai.init(project=PROJECT_ID, location=\"us-east1\")\n",
"client = genai.Client(vertexai=True, project=PROJECT_ID, location=\"global\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "df91814b0f3a"
},
"source": [
"### Import Libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1cc98c707bfc"
},
"outputs": [],
"source": [
"from IPython.display import Markdown, display\n",
"from google.cloud import bigquery\n",
"from google.genai.types import GenerateContentConfig, Retrieval, Tool, VertexRagStore\n",
"from vertexai.preview import rag\n",
"from vertexai.resources.preview import feature_store"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "adbe5c6b3549"
},
"source": [
"## Set up Vertex AI Feature Store\n",
"\n",
"Vertex AI Feature Store, a managed cloud-native service, is an essential component of Vertex AI. It simplifies machine learning (ML) feature management and online serving by letting you manage feature data within a BigQuery table or view. This enables low-latency online feature serving.\n",
"\n",
"For `FeatureOnlineStore` instances created with optimized online serving, you\n",
"can take advantage of a vector similarity search to retrieve a list of\n",
"semantically similar or related entities, which are known as\n",
"*approximate nearest neighbors*.\n",
"\n",
"The following sections show you how to set up a Vertex AI Feature Store instance for your RAG application.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9d54be277ff0"
},
"source": [
"### Create a BigQuery table schema\n",
"\n",
"Use the Cloud Console or the code below to create a BigQuery table schema. It\n",
"must contain the following fields to serve as the data source.\n",
"\n",
"| Field name | Data type | Status |\n",
"|-------------|-----------|--------|\n",
"| `corpus_id` | `String` | Required |\n",
"| `file_id` | `String` | Required |\n",
"| `chunk_id` | `String` | Required |\n",
"| `chunk_data_type` |`String` | Nullable |\n",
"| `chunk_data` | `String` | Nullable |\n",
"| `file_original_uri` | `String` | Nullable |\n",
"| `embeddings` | `Float` | Repeated |\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "c65a3aef9e43"
},
"outputs": [],
"source": [
"bq_client = bigquery.Client(project=PROJECT_ID)\n",
"\n",
"# Define dataset and table name\n",
"dataset_id = \"input_us_central1\" # @param {type:\"string\"}\n",
"table_id = \"rag_source_new\" # @param {type:\"string\"}\n",
"\n",
"schema = [\n",
" bigquery.SchemaField(\"corpus_id\", \"STRING\", mode=\"REQUIRED\"),\n",
" bigquery.SchemaField(\"file_id\", \"STRING\", mode=\"REQUIRED\"),\n",
" bigquery.SchemaField(\"chunk_id\", \"STRING\", mode=\"REQUIRED\"),\n",
" bigquery.SchemaField(\"chunk_data_type\", \"STRING\", mode=\"NULLABLE\"),\n",
" bigquery.SchemaField(\"chunk_data\", \"STRING\", mode=\"NULLABLE\"),\n",
" bigquery.SchemaField(\"file_original_uri\", \"STRING\", mode=\"NULLABLE\"),\n",
" bigquery.SchemaField(\"embeddings\", \"FLOAT64\", mode=\"REPEATED\"),\n",
"]\n",
"\n",
"dataset_ref = bigquery.DatasetReference(PROJECT_ID, dataset_id)\n",
"\n",
"try:\n",
" dataset = bq_client.get_dataset(dataset_ref)\n",
" print(f\"Dataset {dataset_id} already exists.\")\n",
"except Exception:\n",
" dataset = bigquery.Dataset(dataset_ref)\n",
" dataset.location = \"US\" # Set the location (optional, adjust if needed)\n",
" dataset = bq_client.create_dataset(dataset)\n",
" print(f\"Created dataset {dataset.dataset_id}\")\n",
"\n",
"table_ref = dataset_ref.table(table_id)\n",
"table = bq_client.create_table(bigquery.Table(table_ref, schema=schema))\n",
"print(f\"Created table {PROJECT_ID}.{dataset_id}.{table_id}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "f3ed9c71ea1f"
},
"outputs": [],
"source": [
"BIGQUERY_TABLE = f\"bq://{table.full_table_id.replace(':', '.')}\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "93b8a1ce48a9"
},
"source": [
"### Provision a `FeatureOnlineStore` instance\n",
"\n",
"To enable online serving of features, use the `CreateFeatureOnlineStore` API to set up a `FeatureOnlineStore` instance. If you\n",
"are provisioning a `FeatureOnlineStore` for the first time, the operation might\n",
"take approximately five minutes to complete."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ee177c9bc175"
},
"outputs": [],
"source": [
"FEATURE_ONLINE_STORE_ID = \"your_feature_online_store_id\" # @param {type: \"string\"}\n",
"\n",
"fos = feature_store.FeatureOnlineStore.create_optimized_store(FEATURE_ONLINE_STORE_ID)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "07030f075c06"
},
"source": [
"### Create a `FeatureView` resource\n",
"\n",
"To connect the BigQuery table, which stores the feature data source, to\n",
"the `FeatureOnlineStore` instance, call the `CreateFeatureView` API to create a\n",
"`FeatureView` resource. When you create a `FeatureView` resource, choose the\n",
"default distance metric `DOT_PRODUCT_DISTANCE`, which is defined as the\n",
"negative of the dot product (smaller `DOT_PRODUCT_DISTANCE` indicates higher\n",
"similarity)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "013dc5a45561"
},
"outputs": [],
"source": [
"FEATURE_VIEW_ID = \"your_feature_view_id\" # @param {type: \"string\"}\n",
"fv = fos.create_feature_view(\n",
" name=FEATURE_VIEW_ID,\n",
" source=feature_store.utils.FeatureViewVertexRagSource(uri=BIGQUERY_TABLE),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2fdbc1e83aff"
},
"outputs": [],
"source": [
"# Check that Feature View was created\n",
"print(fv)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EdvJRUWRNGHE"
},
"source": [
"## Use Vertex AI Feature Store in RAG Engine\n",
"\n",
"After the Feature Store instance is set up, the following\n",
"sections show you how to set it up as the vector database to use with the RAG\n",
"application."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cd05469c3e71"
},
"source": [
"### Set the vector database to create a RAG corpus\n",
"\n",
"To create the RAG corpus, you must use `FEATURE_VIEW_RESOURCE_NAME`. The\n",
"RAG corpus is created and automatically associated with the\n",
"Vertex AI Feature Store instance.\n",
"\n",
"RAG APIs use the generated `rag_corpus_id` to handle the data upload to the Vertex AI Feature Store\n",
"instance and to retrieve relevant contexts from the `rag_corpus_id`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "53865b3ea33e"
},
"outputs": [],
"source": [
"vector_db = rag.VertexFeatureStore(resource_name=fv.resource_name)\n",
"\n",
"# Name your corpus\n",
"DISPLAY_NAME = \"Feature Store Corpus\" # @param {type:\"string\"}\n",
"\n",
"# Create RAG Corpus\n",
"rag_corpus = rag.create_corpus(display_name=DISPLAY_NAME, vector_db=vector_db)\n",
"print(f\"Created RAG Corpus resource: {rag_corpus.name}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "834931d4b187"
},
"outputs": [],
"source": [
"rag_corpus"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "93a3296647a2"
},
"source": [
"## Import files into the BigQuery table using the RAG API\n",
"\n",
"Use the `ImportRagFiles` API to import files from Google Cloud Storage or\n",
"Google Drive into the BigQuery table of the Vertex AI Feature Store\n",
"instance. The files are embedded and stored in the BigQuery table.\n",
"\n",
"Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Google Cloud Storage bucket"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e0e53a05445e"
},
"outputs": [],
"source": [
"# fmt: off\n",
"GCS_BUCKET = \"cloud-samples-data/gen-app-builder/search/cymbal-bank-employee\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"response = rag.import_files(\n",
" corpus_name=rag_corpus.name,\n",
" paths=[GCS_BUCKET],\n",
" chunk_size=512,\n",
" chunk_overlap=50,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "48313a38ef52"
},
"outputs": [],
"source": [
"# Check the files just imported. It may take a few seconds to process the imported files.\n",
"rag.list_files(corpus_name=rag_corpus.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "b524b1e3fa56"
},
"source": [
"### Run a synchronization process to construct a `FeatureOnlineStore` index {:#run-sync-process}\n",
"\n",
"After uploading your data into the BigQuery table, run a\n",
"synchronization process to make your data available for online serving. You must\n",
"generate a `FeatureOnlineStore` index using the `FeatureView`, and the\n",
"synchronization process might take 20 minutes to complete.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "479d2e1f7725"
},
"outputs": [],
"source": [
"feature_view_sync = fv.sync()\n",
"feature_view_sync"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "63ef904c92b4"
},
"outputs": [],
"source": [
"# Optional: Wait for sync to complete\n",
"feature_view_sync.wait()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "346ceb446e7c"
},
"source": [
"## Use your RAG Corpus to add context to your Gemini queries\n",
"\n",
"When retrieved contexts similarity distance < `vector_distance_threshold`, the contexts (from `RagStore`) will be used for content generation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "fec72ac982c3"
},
"outputs": [],
"source": [
"# Create a tool for the RAG Corpus\n",
"rag_retrieval_tool = Tool(\n",
" retrieval=Retrieval(\n",
" vertex_rag_store=VertexRagStore(\n",
" rag_corpora=[rag_corpus.name],\n",
" similarity_top_k=10,\n",
" vector_distance_threshold=0.4,\n",
" )\n",
" )\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "bff529ef3710"
},
"outputs": [],
"source": [
"MODEL_ID = \"gemini-3.5-flash\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cc0ee39e50f6"
},
"outputs": [],
"source": [
"# fmt: off\n",
"GENERATE_CONTENT_PROMPT = \"What is RAG and why it is helpful?\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"response = client.models.generate_content(\n",
" model=MODEL_ID,\n",
" contents=GENERATE_CONTENT_PROMPT,\n",
" config=GenerateContentConfig(tools=[rag_retrieval_tool]),\n",
")\n",
"\n",
"display(Markdown(response.text))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2899daa12fac"
},
"source": [
"## Using other generation APIs 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": "d549fb12733f"
},
"outputs": [],
"source": [
"RETRIEVAL_QUERY = \"What is RAG and why it is helpful?\" # @param {type:\"string\"}\n",
"\n",
"response = rag.retrieval_query(\n",
" rag_resources=[\n",
" rag.RagResource(\n",
" rag_corpus=rag_corpus.name, # Currently only 1 corpus is allowed.\n",
" )\n",
" ],\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 resources created in this notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ea74a96756a3"
},
"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_feature_store.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
gemini/rag-engine/rag_engine_pinecone.ipynb
+958
View File
@@ -0,0 +1,958 @@
{
"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 Pinecone\n",
"\n",
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_pinecone.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Frag_engine_pinecone.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/rag_engine_pinecone.ipynb\">\n",
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_pinecone.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_pinecone.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_pinecone.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_pinecone.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_pinecone.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_pinecone.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a> "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "84f0f73a0f76"
},
"source": [
"| | |\n",
"|-|-|\n",
"| Author(s) | [Darshan Mehta](https://github.com/darshanmehta17) |"
]
},
{
"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 [Pinecone](https://www.pinecone.io/) as a vector database.\n",
"\n",
"For more information, refer to the [official documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/use-pinecone).\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 --quiet google-cloud-aiplatform google-cloud-secret-manager \"pinecone[grpc]\" google-genai"
]
},
{
"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": [
"<div class=\"alert alert-block alert-warning\">\n",
"<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
"</div>\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dmWOrTJ3gx13"
},
"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": null,
"metadata": {
"id": "NyKGtVQjgx13"
},
"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": "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": null,
"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",
"from google import genai\n",
"\n",
"# fmt: off\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
"# fmt: on\n",
"if not PROJECT_ID or 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)\n",
"client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "adbe5c6b3549"
},
"source": [
"## (Optional) Setup Pinecone Instance\n",
"\n",
"In this section, we have some helper methods to help you setup your Pinecone instance.\n",
"\n",
"Follow the [Pinecone Quickstart](https://docs.pinecone.io/guides/get-started/quickstart) to get an API Key.\n",
"\n",
"This section is not required if you already have a Pinecone instance ready to use."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5303c05f7aa6"
},
"source": [
"### Initialize the Pinecone Python client"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "95f06ed30a65"
},
"outputs": [],
"source": [
"from pinecone import PodSpec, ServerlessSpec\n",
"from pinecone.grpc import PineconeGRPC as Pinecone\n",
"\n",
"# Set API Key\n",
"# Copy this value from your pinecone.io console\n",
"# fmt: off\n",
"PINECONE_API_KEY = \"\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"pc = Pinecone(api_key=PINECONE_API_KEY)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "e821e27f013e"
},
"source": [
"### Create a Pinecone instance\n",
"\n",
"Use the below section to create Pinecone indexes with a spec of your choice (serverless or pod). Read more about Pinecone indexes [here](https://docs.pinecone.io/guides/indexes/create-an-index)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9f521dc70361"
},
"source": [
"#### Serverless Index"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8b465b3d0eca"
},
"outputs": [],
"source": [
"# Index Configs\n",
"INDEX_NAME = \"\" # @param {type:\"string\"}\n",
"\n",
"# Choose a distance metric\n",
"# fmt: off\n",
"DISTANCE_METRIC = \"cosine\" # @param [\"cosine\", \"euclidean\", \"dotproduct\"] {allow-input: true}\n",
"# fmt: on\n",
"\n",
"# This number should match the dimension size of the embedding model you choose\n",
"# for your RAG Corpus.\n",
"EMBEDDING_DIMENSION_SIZE = 768 # @param {\"type\":\"number\",\"placeholder\":\"768\"}\n",
"\n",
"CLOUD_PROVIDER = \"gcp\" # @param [\"gcp\", \"aws\", \"azure\"] {allow-input: true}\n",
"\n",
"# Choose the right region for your cloud provider of choice\n",
"# Refer https://docs.pinecone.io/guides/indexes/understanding-indexes#cloud-regions\n",
"CLOUD_REGION = \"us-central1\" # @param {type:\"string\"}\n",
"\n",
"\n",
"# Create the index\n",
"pc.create_index(\n",
" name=INDEX_NAME,\n",
" dimension=EMBEDDING_DIMENSION_SIZE,\n",
" metric=DISTANCE_METRIC,\n",
" spec=ServerlessSpec(cloud=CLOUD_PROVIDER, region=CLOUD_REGION),\n",
" deletion_protection=\"disabled\",\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "daec49200393"
},
"source": [
"#### Pod-based Index\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "a575d67150b3"
},
"outputs": [],
"source": [
"# Index Configs\n",
"INDEX_NAME = \"\" # @param {type:\"string\"}\n",
"\n",
"# Choose a distance metric\n",
"# fmt: off\n",
"DISTANCE_METRIC = \"cosine\" # @param [\"cosine\", \"euclidean\", \"dotproduct\"] {allow-input: true}\n",
"# fmt: on\n",
"\n",
"# This number should match the dimension size of the embedding model you choose\n",
"# for your RAG Corpus.\n",
"EMBEDDING_DIMENSION_SIZE = 768 # @param {\"type\":\"number\",\"placeholder\":\"768\"}\n",
"\n",
"# Choose the right environment for your cloud provider of choice\n",
"# Refer https://docs.pinecone.io/guides/indexes/understanding-indexes#pod-environments\n",
"ENVIRONMENT = \"us-central1-gcp\" # @param {type:\"string\"}\n",
"\n",
"# Choose the pod type\n",
"# Refer to https://docs.pinecone.io/guides/indexes/understanding-indexes#pod-based-indexes\n",
"POD_TYPE = \"p1.x1\" # @param {type:\"string\"}\n",
"\n",
"# Explore all the parameters you can play with for creating a pod index by\n",
"# following this page:\n",
"# https://docs.pinecone.io/reference/api/2024-07/control-plane/create_index\n",
"pc.create_index(\n",
" name=INDEX_NAME,\n",
" dimension=EMBEDDING_DIMENSION_SIZE,\n",
" metric=DISTANCE_METRIC,\n",
" spec=PodSpec(\n",
" environment=ENVIRONMENT,\n",
" pod_type=POD_TYPE,\n",
" pods=1,\n",
" metadata_config={\n",
" \"indexed\": [\"file_id\"] # This field is required for pod-based indexes.\n",
" },\n",
" ),\n",
" deletion_protection=\"disabled\",\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "968f535385be"
},
"source": [
"## (Optional) Setup Secret Manager\n",
"\n",
"This section helps you add your Pinecone API key to your Google Cloud Secret Manager. This section is not required is you already have a secret with the API key ready to use."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "535dbdfff303"
},
"outputs": [],
"source": [
"# Google Cloud project ID and the Pinecone API key will be used from the above sections.\n",
"# Choose your secret ID\n",
"# fmt: off\n",
"SECRET_ID = \"\" # @param {type:\"string\"}\n",
"# fmt: on"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "217cf0e74015"
},
"outputs": [],
"source": [
"from google.cloud import secretmanager\n",
"\n",
"secretmanager_client = secretmanager.SecretManagerServiceClient()\n",
"\n",
"# Create the secret.\n",
"secret = secretmanager_client.create_secret(\n",
" parent=secretmanager_client.common_project_path(PROJECT_ID),\n",
" secret_id=SECRET_ID,\n",
" secret=secretmanager.Secret(\n",
" replication=secretmanager.Replication(\n",
" automatic=secretmanager.Replication.Automatic()\n",
" )\n",
" ),\n",
")\n",
"\n",
"# Add API key to the secret payload.\n",
"secret_version = secretmanager_client.add_secret_version(\n",
" parent=secret.name,\n",
" payload=secretmanager.SecretPayload(data=PINECONE_API_KEY.encode(\"UTF-8\")),\n",
")\n",
"\n",
"print(f\"Created secret and added first version: {secret_version.name}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9769f99f15de"
},
"source": [
"## Get Service Account Information"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "04e493024500"
},
"outputs": [],
"source": [
"project_numbers = !gcloud projects list --filter=\"PROJECT_ID={PROJECT_ID}\" --format=\"value(PROJECT_NUMBER)\"\n",
"PROJECT_NUMBER = project_numbers[0]\n",
"\n",
"# Construct your RAG Engine service account name\n",
"# Do not update this string since this is the name assigned to your service account.\n",
"SERVICE_ACCOUNT = f\"service-{PROJECT_NUMBER}@gcp-sa-vertex-rag.iam.gserviceaccount.com\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EdvJRUWRNGHE"
},
"source": [
"## Create a RAG corpus"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "b9ad5442bd4e"
},
"outputs": [],
"source": [
"from google.genai.types import (\n",
" GenerateContentConfig,\n",
" Retrieval,\n",
" Tool,\n",
" VertexRagStore,\n",
" VertexRagStoreRagResource,\n",
")\n",
"from vertexai import rag"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "b63b525cee5d"
},
"source": [
"### First RAG Corpus Only\n",
"\n",
"If this is the first RAG Corpus in your project, you might not be able to provide the RAG Engine service account access to your secret resource. So this section first creates a RAG Corpus with an empty Pinecone config. With this call, the service account for your project is provisioned.\n",
"\n",
"Next, it assigns the service account permissions to read your secret. Finally, it updates your RAG Corpus with the Pinecone index name and the secret resource name."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "f46f7f34b390"
},
"source": [
"#### Create a RAG Corpus without Pinecone information"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9a0e7d8bafeb"
},
"outputs": [],
"source": [
"# Start with empty Pinecone config.\n",
"vector_db = rag.Pinecone()\n",
"\n",
"# Name your corpus\n",
"# fmt: off\n",
"DISPLAY_NAME = \"\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"# Create RAG Corpus\n",
"rag_corpus = rag.create_corpus(\n",
" display_name=DISPLAY_NAME, backend_config=rag.RagVectorDbConfig(vector_db=vector_db)\n",
")\n",
"print(f\"Created RAG Corpus resource: {rag_corpus.name}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2422ddb6acd4"
},
"source": [
"#### Grant your RAG Engine service account access to your API key secret"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "493b75f6dee0"
},
"outputs": [],
"source": [
"!gcloud secrets add-iam-policy-binding {secret.name} \\\n",
" --member=\"serviceAccount:{SERVICE_ACCOUNT}\" \\\n",
" --role=\"roles/secretmanager.secretAccessor\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "4e911dd1dcab"
},
"source": [
"#### Call the `UpdateRagCorpus` API to add the Pinecone index name and API key secret to your RAG Corpus"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "84f0c2df5238"
},
"outputs": [],
"source": [
"# Name of your created Pinecone Index\n",
"# fmt: off\n",
"PINECONE_INDEX_NAME = \"\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"# Construct your updated Pinecone config.\n",
"vector_db = rag.Pinecone(index_name=PINECONE_INDEX_NAME, api_key=secret_version.name)\n",
"\n",
"updated_rag_corpora = rag.update_corpus(\n",
" corpus_name=rag_corpus.name,\n",
" backend_config=rag.RagVectorDbConfig(vector_db=vector_db),\n",
")\n",
"print(f\"Updated RAG Corpus: {rag_corpus.name}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "638af13047a1"
},
"source": [
"### Second RAG Corpus Onwards\n",
"\n",
"In this case, since your service account is already generated, you can directly grant it permissions to access your secret resource containing the Pinecone API key as covered by the steps in the Setup Secret Manager section.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "87459cdb45c6"
},
"source": [
"#### Grant your RAG Engine service account access to your API key secret"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "25f9149874fa"
},
"outputs": [],
"source": [
"!gcloud secrets add-iam-policy-binding {secret.name} \\\n",
" --member=\"serviceAccount:{SERVICE_ACCOUNT}\" \\\n",
" --role=\"roles/secretmanager.secretAccessor\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "c1396e74d1cb"
},
"source": [
"#### Create a RAG Corpus with Pinecone information"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "53865b3ea33e"
},
"outputs": [],
"source": [
"# Name of your created Pinecone Index\n",
"# fmt: off\n",
"PINECONE_INDEX_NAME = \"\" # @param {type:\"string\"}\n",
"# Construct your Pinecone config.\n",
"vector_db = rag.Pinecone(index_name=PINECONE_INDEX_NAME, api_key=secret_version.name)\n",
"\n",
"# Name your corpus\n",
"DISPLAY_NAME = \"\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"# Create RAG Corpus\n",
"rag_corpus = rag.create_corpus(\n",
" display_name=DISPLAY_NAME, backend_config=rag.RagVectorDbConfig(vector_db=vector_db)\n",
")\n",
"print(f\"Created RAG Corpus resource: {rag_corpus.name}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "93a3296647a2"
},
"source": [
"## Upload a file to the corpus"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7f31cc83fb04"
},
"outputs": [],
"source": [
"%%writefile test.txt\n",
"\n",
"Here's a demo for Pinecone RAG."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7bab0e824c3d"
},
"outputs": [],
"source": [
"rag_file = rag.upload_file(\n",
" corpus_name=rag_corpus.name,\n",
" path=\"test.txt\",\n",
" display_name=\"test.txt\",\n",
" description=\"my test\",\n",
")\n",
"print(f\"Uploaded file to resource: {rag_file.name}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "e51b5bcd1739"
},
"source": [
"## Import files from Google Cloud Storage\n",
"\n",
"Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Google Cloud Storage bucket"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e0e53a05445e"
},
"outputs": [],
"source": [
"# fmt: off\n",
"GCS_BUCKET = \"\" # @param {type:\"string\", \"placeholder\": \"your-gs-bucket\"}\n",
"# fmt: on\n",
"\n",
"response = rag.import_files(\n",
" corpus_name=rag_corpus.name,\n",
" paths=[GCS_BUCKET],\n",
" transformation_config=rag.TransformationConfig(\n",
" chunking_config=rag.ChunkingConfig(\n",
" chunk_size=512,\n",
" chunk_overlap=50,\n",
" )\n",
" ),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "48313a38ef52"
},
"outputs": [],
"source": [
"# Check the files just imported. It may take a few seconds to process the imported files.\n",
"rag.list_files(corpus_name=rag_corpus.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ceab91983444"
},
"source": [
"## Import files from Google Drive\n",
"\n",
"Eligible paths can be:\n",
"\n",
"- `https://drive.google.com/drive/folders/{folder_id}`\n",
"- `https://drive.google.com/file/d/{file_id}`\n",
"\n",
"Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Drive folder/files.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ea8a5c97ad80"
},
"outputs": [],
"source": [
"# fmt: off\n",
"FILE_ID = \"\" # @param {type:\"string\", \"placeholder\": \"your-file-id\"}\n",
"# fmt: on\n",
"FILE_PATH = f\"https://drive.google.com/file/d/{FILE_ID}\"\n",
"\n",
"rag.import_files(\n",
" corpus_name=rag_corpus.name,\n",
" paths=[FILE_PATH],\n",
" transformation_config=rag.TransformationConfig(\n",
" chunking_config=rag.ChunkingConfig(\n",
" chunk_size=1024,\n",
" chunk_overlap=100,\n",
" )\n",
" ),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e71887752baa"
},
"outputs": [],
"source": [
"# Check the files just imported. It may take a few seconds to process the imported files.\n",
"rag.list_files(corpus_name=rag_corpus.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "346ceb446e7c"
},
"source": [
"## Use your RAG Corpus to add context to your Gemini queries\n",
"\n",
"When retrieved contexts similarity distance < `vector_distance_threshold`, the contexts (from `RagStore`) will be used for content generation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "fec72ac982c3"
},
"outputs": [],
"source": [
"MODEL_ID = \"gemini-2.5-flash\"\n",
"\n",
"rag_retrieval_tool = Tool(\n",
" retrieval=Retrieval(\n",
" vertex_rag_store=VertexRagStore(\n",
" rag_resources=[\n",
" VertexRagStoreRagResource(\n",
" rag_corpus=rag_corpus.name # Currently only 1 corpus is allowed.\n",
" )\n",
" ],\n",
" similarity_top_k=10,\n",
" vector_distance_threshold=0.4,\n",
" )\n",
" )\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cc0ee39e50f6"
},
"outputs": [],
"source": [
"# fmt: off\n",
"GENERATE_CONTENT_PROMPT = \"What is RAG and why it is helpful?\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"response = client.models.generate_content(\n",
" model=MODEL_ID,\n",
" contents=GENERATE_CONTENT_PROMPT,\n",
" config=GenerateContentConfig(tools=[rag_retrieval_tool]),\n",
")\n",
"\n",
"display(Markdown(response.text))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2899daa12fac"
},
"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": "d549fb12733f"
},
"outputs": [],
"source": [
"# fmt: off\n",
"RETRIEVAL_QUERY = \"What is RAG and why it is helpful?\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"rag_resource = rag.RagResource(\n",
" rag_corpus=rag_corpus.name,\n",
" # Need to manually get the ids from rag.list_files.\n",
" # rag_file_ids=[],\n",
")\n",
"\n",
"response = rag.retrieval_query(\n",
" rag_resources=[rag_resource], # Currently only 1 corpus is allowed.\n",
" text=RETRIEVAL_QUERY,\n",
" rag_retrieval_config=rag.RagRetrievalConfig(\n",
" top_k=10, # Optional\n",
" filter=rag.Filter(\n",
" vector_distance_threshold=0.5, # Optional\n",
" ),\n",
" ),\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 resources created in this notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ea74a96756a3"
},
"outputs": [],
"source": [
"# fmt: off\n",
"delete_rag_corpus = False # @param {type:\"boolean\"}\n",
"# fmt: on\n",
"\n",
"if delete_rag_corpus:\n",
" rag.delete_corpus(name=rag_corpus.name)"
]
}
],
"metadata": {
"colab": {
"name": "rag_engine_pinecone.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
gemini/rag-engine/rag_engine_vector_search.ipynb
+700
View File
@@ -0,0 +1,700 @@
{
"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 Vector Search\n",
"\n",
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_vector_search.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Frag_engine_vector_search.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/rag_engine_vector_search.ipynb\">\n",
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_vector_search.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_vector_search.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_vector_search.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_vector_search.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_vector_search.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_vector_search.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a> "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "84f0f73a0f76"
},
"source": [
"| Author |\n",
"| --- |\n",
"| [Holt Skinner](https://github.com/holtskinner) |"
]
},
{
"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 Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview) as a vector database.\n",
"\n",
"For more information, refer to the [official documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/use-vertexai-vector-search).\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": 1,
"metadata": {
"id": "tFy3H3aPgx12"
},
"outputs": [],
"source": [
"%pip install --upgrade --quiet google-cloud-aiplatform google-genai"
]
},
{
"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": [
"<div class=\"alert alert-block alert-warning\">\n",
"<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
"</div>\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dmWOrTJ3gx13"
},
"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": null,
"metadata": {
"id": "NyKGtVQjgx13"
},
"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": "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": null,
"metadata": {
"id": "Nqwi-5ufWp_B"
},
"outputs": [],
"source": [
"# Use the environment variable if the user doesn't provide Project ID.\n",
"import os\n",
"\n",
"from google import genai\n",
"from google.cloud import aiplatform\n",
"\n",
"# fmt: off\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
"# fmt: on\n",
"if not PROJECT_ID or 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",
"aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
"client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "adbe5c6b3549"
},
"source": [
"## (Optional) Setup Vertex AI Vector Search index and index endpoint\n",
"\n",
"In this section, we have some helper methods to help you setup your Vector Search index.\n",
"\n",
"This section is not required if you already have a Vector Search index ready to use.\n",
"\n",
"The index has to meet the following criteria:\n",
"\n",
"1. `IndexUpdateMethod` must be `STREAM_UPDATE`, see [Create stream index]({{docs_path}}vector-search/create-manage-index#create-stream-index).\n",
"\n",
"2. Distance measure type must be explicitly set to one of the following:\n",
"\n",
" * `DOT_PRODUCT_DISTANCE`\n",
" * `COSINE_DISTANCE`\n",
"\n",
"3. Dimension of the vector must be consistent with the embedding model you plan\n",
" to use in the RAG corpus. Other parameters can be tuned based on\n",
" your choices, which determine whether the additional parameters can be\n",
" tuned."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ee177c9bc175"
},
"outputs": [],
"source": [
"# create the index\n",
"my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
" display_name=\"your_display_name\",\n",
" description=\"your_description\",\n",
" dimensions=768,\n",
" approximate_neighbors_count=10,\n",
" leaf_node_embedding_count=500,\n",
" leaf_nodes_to_search_percent=7,\n",
" distance_measure_type=\"DOT_PRODUCT_DISTANCE\",\n",
" feature_norm_type=\"UNIT_L2_NORM\",\n",
" index_update_method=\"STREAM_UPDATE\",\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "02e52d8dcda6"
},
"source": [
"RAG Engine supports [public endpoints](https://cloud.google.com/vertex-ai/docs/vector-search/deploy-index-public)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ce6e0e85adf1"
},
"outputs": [],
"source": [
"# create IndexEndpoint\n",
"my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
" display_name=\"your_display_name\", public_endpoint_enabled=True\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "c4c3f91ab95f"
},
"source": [
"Deploy the index to the index endpoint.\n",
"\n",
"If it's the first time that you're deploying an index to an index endpoint, it\n",
"takes approximately 30 minutes to automatically build and initiate the backend\n",
"before the index can be stored. After the first deployment, the index is ready\n",
"in seconds. To see the status of the index deployment, open the\n",
"[**Vector Search Console**](https://console.cloud.google.com/vertex-ai/matching-engine/index-endpoints),\n",
"select the **Index endpoints** tab, and choose your index endpoint.\n",
"\n",
"Identify the resource name of your index and index endpoint, which have the\n",
"following the formats:\n",
"\n",
"* `projects/${PROJECT_ID}/locations/${LOCATION_ID}/indexes/${INDEX_ID}`\n",
"* `projects/${PROJECT_ID}/locations/${LOCATION_ID}/indexEndpoints/${INDEX_ENDPOINT_ID}`.\n",
"\n",
"If you aren't sure about the resource name, you can use the following command to\n",
"check:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "382010f08560"
},
"outputs": [],
"source": [
"print(my_index_endpoint.resource_name)\n",
"print(my_index.resource_name)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6dab214cd107"
},
"outputs": [],
"source": [
"# Deploy Index\n",
"my_index_endpoint.deploy_index(\n",
" index=my_index, deployed_index_id=\"your_deployed_index_id\"\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EdvJRUWRNGHE"
},
"source": [
"## Use Vertex AI Vector Search in RAG Engine\n",
"\n",
"After the Vector Search instance is set up, follow the steps in this section to set the Vector Search instance as the vector database for the RAG application.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cd05469c3e71"
},
"source": [
"### Set the vector database to create a RAG corpus"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "b9ad5442bd4e"
},
"outputs": [],
"source": [
"from google.genai.types import (\n",
" GenerateContentConfig,\n",
" Retrieval,\n",
" Tool,\n",
" VertexRagStore,\n",
" VertexRagStoreRagResource,\n",
")\n",
"from vertexai import rag"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "53865b3ea33e"
},
"outputs": [],
"source": [
"vector_db = rag.VertexVectorSearch(\n",
" index=my_index.resource_name, index_endpoint=my_index_endpoint.resource_name\n",
")\n",
"\n",
"# Name your corpus\n",
"DISPLAY_NAME = \"\" # @param {type:\"string\"}\n",
"\n",
"# Create RAG Corpus\n",
"rag_corpus = rag.create_corpus(\n",
" display_name=DISPLAY_NAME, backend_config=rag.RagVectorDbConfig(vector_db=vector_db)\n",
")\n",
"print(f\"Created RAG Corpus resource: {rag_corpus.name}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "93a3296647a2"
},
"source": [
"## Upload a file to the corpus"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7f31cc83fb04"
},
"outputs": [],
"source": [
"%%writefile test.txt\n",
"\n",
"Here's a demo for Vertex AI Vector Search RAG."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7bab0e824c3d"
},
"outputs": [],
"source": [
"rag_file = rag.upload_file(\n",
" corpus_name=rag_corpus.name,\n",
" path=\"test.txt\",\n",
" display_name=\"test.txt\",\n",
" description=\"my test\",\n",
")\n",
"print(f\"Uploaded file to resource: {rag_file.name}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "e51b5bcd1739"
},
"source": [
"## Import files from Google Cloud Storage\n",
"\n",
"Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Google Cloud Storage bucket"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e0e53a05445e"
},
"outputs": [],
"source": [
"GCS_BUCKET = \"\" # @param {type:\"string\", \"placeholder\": \"your-gs-bucket\"}\n",
"\n",
"response = rag.import_files(\n",
" corpus_name=rag_corpus.name,\n",
" paths=[GCS_BUCKET],\n",
" transformation_config=rag.TransformationConfig(\n",
" chunking_config=rag.ChunkingConfig(\n",
" chunk_size=512,\n",
" chunk_overlap=50,\n",
" )\n",
" ),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "48313a38ef52"
},
"outputs": [],
"source": [
"# Check the files just imported. It may take a few seconds to process the imported files.\n",
"rag.list_files(corpus_name=rag_corpus.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ceab91983444"
},
"source": [
"## Import files from Google Drive\n",
"\n",
"Eligible paths can be:\n",
"\n",
"- `https://drive.google.com/drive/folders/{folder_id}`\n",
"- `https://drive.google.com/file/d/{file_id}`\n",
"\n",
"Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Drive folder/files.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ea8a5c97ad80"
},
"outputs": [],
"source": [
"FILE_ID = \"\" # @param {type:\"string\", \"placeholder\": \"your-file-id\"}\n",
"FILE_PATH = f\"https://drive.google.com/file/d/{FILE_ID}\"\n",
"\n",
"rag.import_files(\n",
" corpus_name=rag_corpus.name,\n",
" paths=[FILE_PATH],\n",
" transformation_config=rag.TransformationConfig(\n",
" chunking_config=rag.ChunkingConfig(\n",
" chunk_size=1024,\n",
" chunk_overlap=100,\n",
" )\n",
" ),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e71887752baa"
},
"outputs": [],
"source": [
"# Check the files just imported. It may take a few seconds to process the imported files.\n",
"rag.list_files(corpus_name=rag_corpus.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "346ceb446e7c"
},
"source": [
"## Use your RAG Corpus to add context to your Gemini queries\n",
"\n",
"When retrieved contexts similarity distance < `vector_distance_threshold`, the contexts (from `RagStore`) will be used for content generation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "fec72ac982c3"
},
"outputs": [],
"source": [
"MODEL_ID = \"gemini-2.5-flash\"\n",
"\n",
"rag_retrieval_tool = Tool(\n",
" retrieval=Retrieval(\n",
" vertex_rag_store=VertexRagStore(\n",
" rag_resources=[\n",
" VertexRagStoreRagResource(\n",
" rag_corpus=rag_corpus.name # Currently only 1 corpus is allowed.\n",
" )\n",
" ],\n",
" similarity_top_k=10,\n",
" vector_distance_threshold=0.4,\n",
" )\n",
" )\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cc0ee39e50f6"
},
"outputs": [],
"source": [
"# fmt: off\n",
"GENERATE_CONTENT_PROMPT = \"What is RAG and why it is helpful?\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"response = client.models.generate_content(\n",
" model=MODEL_ID,\n",
" contents=GENERATE_CONTENT_PROMPT,\n",
" config=GenerateContentConfig(tools=[rag_retrieval_tool]),\n",
")\n",
"\n",
"display(Markdown(response.text))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2899daa12fac"
},
"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": "d549fb12733f"
},
"outputs": [],
"source": [
"RETRIEVAL_QUERY = \"What is RAG and why it is helpful?\" # @param {type:\"string\"}\n",
"\n",
"rag_resource = rag.RagResource(\n",
" rag_corpus=rag_corpus.name,\n",
" # Need to manually get the ids from rag.list_files.\n",
" # rag_file_ids=[],\n",
")\n",
"\n",
"response = rag.retrieval_query(\n",
" rag_resources=[rag_resource], # Currently only 1 corpus is allowed.\n",
" text=RETRIEVAL_QUERY,\n",
" rag_retrieval_config=rag.RagRetrievalConfig(\n",
" top_k=10, # Optional\n",
" filter=rag.Filter(\n",
" vector_distance_threshold=0.5, # Optional\n",
" ),\n",
" ),\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 resources created in this notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ea74a96756a3"
},
"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_vector_search.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
gemini/rag-engine/rag_engine_vertex_ai_search.ipynb
+682
View File
@@ -0,0 +1,682 @@
{
"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",
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_vertex_ai_search.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Frag_engine_vertex_ai_search.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/rag_engine_vertex_ai_search.ipynb\">\n",
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_vertex_ai_search.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</table>"
]
},
{
"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": [
"<div class=\"alert alert-block alert-warning\">\n",
"<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
"</div>\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",
"# fmt: off\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\", isTemplate: true}\n",
"# fmt: on\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",
"# fmt: off\n",
"DATASTORE_NAME = \"alphabet-contracts\" # @param {type:\"string\", isTemplate: true}\n",
"# fmt: on\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": [
"# fmt: off\n",
"GCS_BUCKET = \"gs://cloud-samples-data/gen-app-builder/search/alphabet-investor-pdfs\" # @param {type:\"string\", isTemplate: true}\n",
"# fmt: on\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",
"# fmt: off\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",
"# fmt: on\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-2.0-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": [
"# fmt: off\n",
"GENERATE_CONTENT_PROMPT = \"Who is CFO of Google?\" # @param {type:\"string\", isTemplate: true}\n",
"# fmt: on\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": [
"# fmt: off\n",
"RETRIEVAL_QUERY = \"Who is CFO of Google?\" # @param {type:\"string\", isTemplate: true}\n",
"# fmt: on\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": [
"# fmt: off\n",
"delete_rag_corpus = False # @param {type:\"boolean\"}\n",
"# fmt: on\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
}
gemini/rag-engine/rag_engine_weaviate.ipynb
+594
View File
@@ -0,0 +1,594 @@
{
"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 Weaviate\n",
"\n",
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_weaviate.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Frag_engine_weaviate.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/rag_engine_weaviate.ipynb\">\n",
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_weaviate.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_weaviate.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_weaviate.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_weaviate.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_weaviate.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_weaviate.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a> "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "84f0f73a0f76"
},
"source": [
"| | |\n",
"|-|-|\n",
"| Author(s) | [Ming Zhang](https://github.com/mzhang-ai) |"
]
},
{
"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 [Weaviate](https://weaviate.io/) as a vector database.\n",
"\n",
"For more information, refer to the [official documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/use-weaviate-db).\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"
]
},
{
"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": [
"<div class=\"alert alert-block alert-warning\">\n",
"<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
"</div>\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dmWOrTJ3gx13"
},
"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": null,
"metadata": {
"id": "NyKGtVQjgx13"
},
"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": "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": 3,
"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",
"# fmt: off\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\", isTemplate: true}\n",
"# fmt: on\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": "EdvJRUWRNGHE"
},
"source": [
"## Create a RAG corpus using Weaviate as the Vector Database"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5303c05f7aa6"
},
"source": [
"### Import libraries"
]
},
{
"cell_type": "code",
"execution_count": 4,
"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": [
"### Load embedding model and create RAG Config"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"id": "cf93d5f0ce00"
},
"outputs": [],
"source": [
"# Configure a Google first-party embedding model\n",
"embedding_model_config = rag.EmbeddingModelConfig(\n",
" publisher_model=\"publishers/google/models/text-embedding-005\"\n",
")\n",
"\n",
"# Name your corpus\n",
"DISPLAY_NAME = \"\" # @param {type:\"string\", \"placeholder\": \"your-corpus-name\"}\n",
"\n",
"# Configure a Weaviate Vector Database Instance for the corpus\n",
"# More details for how to deploy a Weaviate Database Instance\n",
"# https://cloud.google.com/vertex-ai/generative-ai/docs/use-weaviate-db\n",
"# fmt: off\n",
"WEAVIATE_HTTP_ENDPOINT = \"\" # @param {type:\"string\", \"placeholder\": \"your-weaviate-http-endpoint\"}\n",
"COLLECTION_NAME = \"\" # @param {type:\"string\", \"placeholder\": \"your-weaviate-collection-name\"}\n",
"API_KEY = \"\" # @param {type:\"string\", \"placeholder\": \"your-secret-manager-resource-name\"}\n",
"# fmt: on\n",
"vector_db = rag.Weaviate(\n",
" weaviate_http_endpoint=WEAVIATE_HTTP_ENDPOINT,\n",
" collection_name=COLLECTION_NAME,\n",
" api_key=API_KEY,\n",
")\n",
"\n",
"rag_corpus = rag.create_corpus(\n",
" display_name=DISPLAY_NAME,\n",
" embedding_model_config=embedding_model_config,\n",
" vector_db=vector_db,\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": "f5f3a12a95ca"
},
"source": [
"## Upload a file to the corpus"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e90f8ddfb7ee"
},
"outputs": [],
"source": [
"%%writefile test.txt\n",
"\n",
"Here's a demo for Weaviate RAG."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "94d84155accc"
},
"outputs": [],
"source": [
"rag_file = rag.upload_file(\n",
" corpus_name=rag_corpus.name,\n",
" path=\"test.txt\",\n",
" display_name=\"test.txt\",\n",
" description=\"my test\",\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2627b2ac8caf"
},
"source": [
"## Import files from Google Cloud Storage\n",
"\n",
"Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Google Cloud Storage bucket"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1c7d989d75bb"
},
"outputs": [],
"source": [
"# fmt: off\n",
"GCS_BUCKET = \"\" # @param {type:\"string\", \"placeholder\": \"your-gs-bucket\"}\n",
"# fmt: on\n",
"\n",
"response = rag.import_files(\n",
" corpus_name=rag_corpus.name,\n",
" paths=[GCS_BUCKET],\n",
" chunk_size=512,\n",
" chunk_overlap=50,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "66c8b9325082"
},
"outputs": [],
"source": [
"# Check the files just imported. It may take a few seconds to process the imported files.\n",
"rag.list_files(corpus_name=rag_corpus.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "43d6712c4bd2"
},
"source": [
"## Import files from Google Drive\n",
"\n",
"Eligible paths can be:\n",
"\n",
"- `https://drive.google.com/drive/folders/{folder_id}`\n",
"- `https://drive.google.com/file/d/{file_id}`\n",
"\n",
"Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Drive folder/files.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "d3df465a3eb2"
},
"outputs": [],
"source": [
"# fmt: off\n",
"FILE_ID = \"\" # @param {type:\"string\", \"placeholder\": \"your-file-id\"}\n",
"# fmt: on\n",
"FILE_PATH = f\"https://drive.google.com/file/d/{FILE_ID}\"\n",
"\n",
"rag.import_files(\n",
" corpus_name=rag_corpus.name,\n",
" paths=[FILE_PATH],\n",
" chunk_size=1024,\n",
" chunk_overlap=100,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "bcfca523ca38"
},
"outputs": [],
"source": [
"# Check the files just imported. It may take a few seconds to process the imported files.\n",
"rag.list_files(corpus_name=rag_corpus.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "e1c1c7944dc8"
},
"source": [
"## Using Gemini GenerateContent API with Rag Retrieval Tool\n",
"\n",
"When retrieved contexts similarity distance < `vector_distance_threshold`, the contexts (from `RagStore`) will be used for content generation."
]
},
{
"cell_type": "code",
"execution_count": null,
"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",
" vector_distance_threshold=0.4,\n",
" ),\n",
" )\n",
")\n",
"\n",
"rag_model = GenerativeModel(\"gemini-2.0-flash\", tools=[rag_retrieval_tool])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "484f5242dcae"
},
"outputs": [],
"source": [
"# fmt: off\n",
"GENERATE_CONTENT_PROMPT = \"What is RAG and why it is helpful?\" # @param {type:\"string\"}\n",
"# fmt: on\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": [
"# fmt: off\n",
"RETRIEVAL_QUERY = \"What is RAG and why it is helpful?\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"rag_resource = rag.RagResource(\n",
" rag_corpus=rag_corpus.name,\n",
" # Need to manually get the ids from rag.list_files.\n",
" # rag_file_ids=[],\n",
")\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 resources created in this notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "a105caffd9e7"
},
"outputs": [],
"source": [
"# fmt: off\n",
"delete_rag_corpus = False # @param {type:\"boolean\"}\n",
"# fmt: on\n",
"\n",
"if delete_rag_corpus:\n",
" rag.delete_corpus(name=rag_corpus.name)"
]
}
],
"metadata": {
"colab": {
"name": "rag_engine_weaviate.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}