gemini/rag-engine/rag_engine_weaviate.ipynb

{
  "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://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_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/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_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",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\", isTemplate: true}\n",
        "if PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "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-004\"\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",
        "WEAVIATE_HTTP_ENDPOINT = (\n",
        "    \"\"  # @param {type:\"string\", \"placeholder\": \"your-weaviate-http-endpoint\"}\n",
        ")\n",
        "COLLECTION_NAME = (\n",
        "    \"\"  # @param {type:\"string\", \"placeholder\": \"your-weaviate-collection-namet\"}\n",
        ")\n",
        "API_KEY = (\n",
        "    \"\"  # @param {type:\"string\", \"placeholder\": \"your-secret-manager-resource-name\"}\n",
        ")\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": [
        "GCS_BUCKET = \"\"  # @param {type:\"string\", \"placeholder\": \"your-gs-bucket\"}\n",
        "\n",
        "response = rag.import_files(  # noqa: F704\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": [
        "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",
        "    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-1.5-flash\", tools=[rag_retrieval_tool])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "484f5242dcae"
      },
      "outputs": [],
      "source": [
        "GENERATE_CONTENT_PROMPT = \"What is RAG and why it is helpful?\"  # @param {type:\"string\"}\n",
        "\n",
        "response = rag_model.generate_content(GENERATE_CONTENT_PROMPT)\n",
        "\n",
        "response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "287a90fed14f"
      },
      "source": [
        "## Using other generation API with Rag Retrieval Tool\n",
        "\n",
        "The retrieved contexts can be passed to any SDK or model generation API to generate final results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "428921dea97d"
      },
      "outputs": [],
      "source": [
        "RETRIEVAL_QUERY = \"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",
        "    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": [
        "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_weaviate.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}