wehub-resource-sync
701 lines
24 KiB
Plaintext
701 lines
24 KiB
Plaintext
{
|
|
"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
|
|
}
|