{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RcDLkWLmmzlB"
      },
      "source": [
        "# Leverage LlamaIndex with Vertex AI Vector Search to perform question answering RAG\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/use-cases/retrieval-augmented-generation/llamaindex_rag.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%2Fuse-cases%2Fretrieval-augmented-generation%2Fllamaindex_rag.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/use-cases/retrieval-augmented-generation/llamaindex_rag.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/use-cases/retrieval-augmented-generation/llamaindex_rag.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/use-cases/retrieval-augmented-generation/llamaindex_rag.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/use-cases/retrieval-augmented-generation/llamaindex_rag.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/use-cases/retrieval-augmented-generation/llamaindex_rag.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/use-cases/retrieval-augmented-generation/llamaindex_rag.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/use-cases/retrieval-augmented-generation/llamaindex_rag.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": "1bc291e87bd1"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) | [Noa Ben-Efraim](https://github.com/noabenefraim/) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ecb6e86acd7f"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook will go over how to create a RAG framework using LlamaIndex and Vertex AI Vector Search effectively.\n",
        "\n",
        "LlamaIndex is used to parse, chunk, and embed the input data using Gemini Text Embedding models. We use store the parsed data in a Vertex AI Vector Search index that will searched against during inference to retrieve context to augment prompts for question answering task.\n",
        "\n",
        "### Objectives\n",
        "This notebook provides a guide to building a questions answering system using retrieval augmented generation (RAG) framework that leverages LlamaIndex for data ingestion and Vector Store creation.\n",
        "\n",
        "You will complete the following tasks:\n",
        "\n",
        "1. Set up Google Cloud resources required: GCS Bucket and Vertex AI Vector Search index and deployed endpoint\n",
        "2. Ingest, parse, chunk, and embed data using LlamaIndex with Gemini Text Embedding models.\n",
        "3. Search the vector store with an incoming text queries to find similar text data that can be used as context in the prompt\n",
        "4. Generate answer to the user query using Gemini Model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vw__bUsqmzlG"
      },
      "source": [
        "### LlamaIndex\n",
        "\n",
        "LlamaIndex is a data framework specifically designed to enhance the capabilities of Large Language Models (LLMs) like Gemini-Pro. It addresses the limitations of LLMs in handling and retrieving information from private or domain-specific data sources.\n",
        "\n",
        "LlamaIndex enables seamless integration with diverse data sources, and enhances the ability to retrieve and utilize relevant information. Therefore it is a great candidate to build RAG frameworks on top of.\n",
        "\n",
        "Key reasons to use LlamaIndex:\n",
        "\n",
        "1. Efficient Data Ingestion and Indexing:\n",
        "\n",
        "    Diverse Data Sources: LlamaIndex simplifies the process of loading data from various sources, including APIs, PDFs, SQL databases, and more. It provides a standardized way to ingest and transform data into a format suitable for LLMs.\n",
        "\n",
        "    Structured Indexing: LlamaIndex creates structured indexes (vector indexes, keyword indexes, etc.) over the ingested data. This allows LLMs to efficiently retrieve relevant information based on queries or prompts.\n",
        "\n",
        "2. Enhanced Retrieval and Contextual Understanding:\n",
        "\n",
        "    Query Engine: LlamaIndex offers a flexible query engine that enables LLMs to search and retrieve information from indexed data based on natural language queries. It can combine keyword-based search with semantic understanding for more accurate results.\n",
        "    Contextual Augmentation: LlamaIndex can augment LLMs with relevant context from the indexed data, improving the quality and accuracy of generated responses.\n",
        "\n",
        "3. Streamlined Development of LLM Applications:\n",
        "\n",
        "    Simplified Integration: LlamaIndex simplifies the integration of LLMs with external data sources, reducing the complexity of building custom LLM applications.\n",
        "\n",
        "    Abstraction of Complexities: It abstracts away many of the technical challenges associated with working with LLMs, such as data preprocessing, indexing, and retrieval, allowing developers to focus on the core application logic.\n",
        "    \n",
        "    Rapid Prototyping: LlamaIndex facilitates rapid prototyping and experimentation with different data sources and LLM configurations, accelerating the development cycle.\n",
        "\n",
        "4. Extensibility and Customization:\n",
        "\n",
        "    Modular Design: LlamaIndex is designed with modularity in mind, allowing you to customize and extend its functionality to suit your specific use case.\n",
        "\n",
        "    Community-Driven Extensions: It has a growing community of developers contributing new data connectors, index types, and query strategies.\n",
        "\n",
        "\n",
        "### Vertex AI\n",
        "This notebook utilizes two main components of Vertex AI:\n",
        "1. Gemini models\n",
        "2. Vertex AI Vector Search\n",
        "\n",
        "The Gemini models used are the Text Embedding model and Gemini-Pro. The embedding model is used to convert textual data into numerical vectors that can be stored and searched against. Gemini-Pro is used to generate the completion using the RAG prompt as input and returning an answer to the query in the specified format.\n",
        "\n",
        "Vertex AI Vector Search is used to search for semantically similar or relevant items. Semantic matching can be simplified into a few steps. To utilize Vector Search, first you must generate embeddings then upload and link your data to Vector Search. Then you can create an index to run queries against to get recommendations or results.\n",
        "\n",
        "### RAG\n",
        "Retrieval augmented generation (RAG) has become a popular paradigm for enabling LLMs to access external data and also as a mechanism for grounding to mitigate against hallucinations.\n",
        "\n",
        "In this notebook, you will learn how to perform RAG where you will perform Q&A over a document and explore how different prompt formats impact RAG results.\n",
        "\n",
        "\n",
        "### Costs\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "+ Vertex AI\n",
        "+ Google Cloud Storage Bucket\n",
        "\n",
        "Learn about Vertex AI pricing and use the Pricing Calculator to generate a cost estimate based on your projected usage"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RUCjnjIymzlH"
      },
      "source": [
        "## Getting Started\n",
        "\n",
        "### Authenticate your notebook enviorment\n",
        "\n",
        "This notebook expects the following resources to exist:\n",
        "+ initialized Google Cloud project \n",
        "+ Vertex AI API enabled\n",
        "+ existing VPC/Subnet \n",
        "\n",
        "Note: _If you have an existing GCS bucket and Vector Search index and endpoint, please update the variables below to correspond to the existing ones._"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "JwgeM75nmzlH"
      },
      "outputs": [],
      "source": [
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # TODO add your project id here # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "REGION = \"us-central1\"  # TODO add your project region here # @param {type:\"string\"}\n",
        "GCS_BUCKET = \"[your-gcs-bucket]\"  # @param {type:\"string\"}\n",
        "VS_INDEX_NAME = \"llamaindex_doc_index\"  # @param {type:\"string\"}\n",
        "VS_INDEX_ENDPOINT_NAME = \"llamaindex_doc_endpoint\"  # @param {type:\"string\"}\n",
        "DOC_FOLDER = \"./data\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e2f2ac96afd1"
      },
      "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": 2,
      "metadata": {
        "id": "324e825fe3bf"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=\"us-central1\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NXxJry6KmzlH"
      },
      "source": [
        "### Setting up the Environment\n",
        "Install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "n-HRXI8enntA"
      },
      "outputs": [],
      "source": [
        "%pip install google-cloud-aiplatform \\\n",
        "  google-cloud-storage \\\n",
        "  llama-index \\\n",
        "  llama-index-embeddings-vertex \\\n",
        "  llama-index-llms-vertex \\\n",
        "  llama-index-vector_stores-vertexaivectorsearch \\\n",
        "  langchain-community \\\n",
        "  llama-index-llms-langchain \\\n",
        "  llama-index-llms-fireworks \\\n",
        "  langchainhub -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l96WnMIow7d-"
      },
      "source": [
        "Set up imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "TH2aWgUgmzlH"
      },
      "outputs": [],
      "source": [
        "# Imports\n",
        "import os\n",
        "\n",
        "from google.cloud import aiplatform, storage\n",
        "from langchain import hub\n",
        "from llama_index.core import (\n",
        "    Document,\n",
        "    PromptTemplate,\n",
        "    Settings,\n",
        "    SimpleDirectoryReader,\n",
        "    StorageContext,\n",
        "    SummaryIndex,\n",
        "    VectorStoreIndex,\n",
        ")\n",
        "from llama_index.core.agent import ReActAgent\n",
        "from llama_index.core.base.base_query_engine import BaseQueryEngine\n",
        "from llama_index.core.node_parser import SentenceSplitter\n",
        "from llama_index.core.objects import ObjectIndex\n",
        "from llama_index.core.prompts import LangchainPromptTemplate\n",
        "from llama_index.core.prompts.base import BasePromptTemplate\n",
        "from llama_index.core.tools import QueryEngineTool, ToolMetadata\n",
        "from llama_index.embeddings.vertex import VertexTextEmbedding\n",
        "from llama_index.llms.vertex import Vertex\n",
        "from llama_index.vector_stores.vertexaivectorsearch import VertexAIVectorStore"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "vJ9N-x2u9OyU"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Updated property [core/project].\n"
          ]
        }
      ],
      "source": [
        "!gcloud config set project {PROJECT_ID}\n",
        "# !gcloud auth application-default login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e48fbda5ff18"
      },
      "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": "7a918c599327"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bef4a1021f6f"
      },
      "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": "f460f33d2158"
      },
      "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": "4x10VEywwN3Z"
      },
      "source": [
        "### Using LlamaIndex"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-2mDLUbdweJi"
      },
      "source": [
        "![LlamaRAG](https://storage.googleapis.com/github-repo/generative-ai/gemini/use-cases/retrieval-augmented-generation/llamaindex/LlamaRAG.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ybyxRabnmzlI"
      },
      "source": [
        "### Download sample data\n",
        "Refer to document 04a02.pdf\n",
        "\n",
        "This document describes the importance of stable power grids in Japan, highlighting the recent failure of a generator step-up transformer at the Nakoso Power Station and the rapid restoration response undertaken to maintain power supply stability.\n",
        "\n",
        "We will use this pdf moving forward."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0f2f9cac2beb"
      },
      "outputs": [],
      "source": [
        "!mkdir {DOC_FOLDER}\n",
        "!gcloud storage cp gs://github-repo/generative-ai/gemini/use-cases/retrieval-augmented-generation/llamaindex/data/* {DOC_FOLDER}\n",
        "\n",
        "print(\"Download completed\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zn2zpsSQmzlJ"
      },
      "source": [
        "## Set Up: Vertex AI Vector Search Index and Endpoint\n",
        "\n",
        "This section goes over how to set up the necessary cloud resources to run the RAG framework. You will be performing the following steps:\n",
        "\n",
        "1. Creating a Google Cloud bucket\n",
        "2. Creating a Vertex AI Vector Search Index\n",
        "3. Creating a Vertex AI Vector Search endpoint\n",
        "4. Deploying a Vertex AI Vector Search endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "f7cD_tHsmzlJ"
      },
      "outputs": [],
      "source": [
        "def create_bucket_class_location(bucket_name: str) -> storage.Bucket:\n",
        "    \"\"\"Create a new bucket in the US region with the coldline storage\n",
        "    class.\n",
        "    \"\"\"\n",
        "    storage_client = storage.Client()\n",
        "\n",
        "    # Searching for existing GCS bucket\n",
        "    for bucket in storage_client.list_buckets():\n",
        "        if bucket.name == bucket_name:\n",
        "            print(f\"GCS Bucket {bucket_name} exists already in resource.\")\n",
        "            return bucket\n",
        "\n",
        "    # Creating new bucket\n",
        "    bucket = storage_client.bucket(bucket_name)\n",
        "    bucket.storage_class = \"STANDARD\"\n",
        "    new_bucket = storage_client.create_bucket(bucket, location=REGION)\n",
        "\n",
        "    print(\n",
        "        f\"Created bucket {new_bucket.name} in {new_bucket.location} with storage class {new_bucket.storage_class}\"\n",
        "    )\n",
        "\n",
        "    return new_bucket\n",
        "\n",
        "\n",
        "def create_vector_search_index(\n",
        "    index_name: str, index_dimensions: int\n",
        ") -> aiplatform.MatchingEngineIndex:\n",
        "    \"\"\"Creates a Vector Index\n",
        "    NOTE : This operation can take upto 30 minutes\n",
        "    \"\"\"\n",
        "    # check if index exists\n",
        "    index_names = [\n",
        "        index.resource_name\n",
        "        for index in aiplatform.MatchingEngineIndex.list(\n",
        "            filter=f\"display_name={index_name}\"\n",
        "        )\n",
        "    ]\n",
        "\n",
        "    if len(index_names) == 0:\n",
        "        print(f\"Creating Vector Search index {index_name} ...\")\n",
        "        vs_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
        "            display_name=index_name,\n",
        "            dimensions=index_dimensions,\n",
        "            # distance_measure_type=\"DOT_PRODUCT_DISTANCE\",\n",
        "            shard_size=\"SHARD_SIZE_SMALL\",\n",
        "            index_update_method=\"STREAM_UPDATE\",  # allowed values BATCH_UPDATE , STREAM_UPDATE,\n",
        "            approximate_neighbors_count=5,\n",
        "        )\n",
        "        print(\n",
        "            f\"Vector Search index {vs_index.display_name} created with resource name {vs_index.resource_name}\"\n",
        "        )\n",
        "    else:\n",
        "        vs_index = aiplatform.MatchingEngineIndex(index_name=index_names[0])\n",
        "        print(\n",
        "            f\"Vector Search index {vs_index.display_name} exists with resource name {vs_index.resource_name}\"\n",
        "        )\n",
        "\n",
        "    return vs_index\n",
        "\n",
        "\n",
        "def create_vector_search_endpoint(\n",
        "    endpoint_name: str,\n",
        ") -> aiplatform.MatchingEngineIndexEndpoint:\n",
        "    \"\"\"Creates a Vector Search endpoint.\"\"\"\n",
        "    endpoint_names = [\n",
        "        endpoint.resource_name\n",
        "        for endpoint in aiplatform.MatchingEngineIndexEndpoint.list(\n",
        "            filter=f\"display_name={endpoint_name}\"\n",
        "        )\n",
        "    ]\n",
        "\n",
        "    if len(endpoint_names) == 0:\n",
        "        print(f\"Creating Vector Search index endpoint {endpoint_name} ...\")\n",
        "        vs_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
        "            display_name=endpoint_name, public_endpoint_enabled=True\n",
        "        )\n",
        "        print(\n",
        "            f\"Vector Search index endpoint {vs_endpoint.display_name} created with resource name {vs_endpoint.resource_name}\"\n",
        "        )\n",
        "    else:\n",
        "        vs_endpoint = aiplatform.MatchingEngineIndexEndpoint(\n",
        "            index_endpoint_name=endpoint_names[0]\n",
        "        )\n",
        "        print(\n",
        "            f\"Vector Search index endpoint {vs_endpoint.display_name} exists with resource name {vs_endpoint.resource_name}\"\n",
        "        )\n",
        "\n",
        "    return vs_endpoint\n",
        "\n",
        "\n",
        "def deploy_vector_search_endpoint(\n",
        "    vs_index: aiplatform.MatchingEngineIndex,\n",
        "    vs_endpoint: aiplatform.MatchingEngineIndexEndpoint,\n",
        "    index_name: str,\n",
        ") -> aiplatform.MatchingEngineIndexEndpoint:\n",
        "    \"\"\"Deploys a Vector Search endpoint.\"\"\"\n",
        "    # check if endpoint exists\n",
        "    index_endpoints = [\n",
        "        (deployed_index.index_endpoint, deployed_index.deployed_index_id)\n",
        "        for deployed_index in vs_index.deployed_indexes\n",
        "    ]\n",
        "\n",
        "    if len(index_endpoints) == 0:\n",
        "        print(\n",
        "            f\"Deploying Vector Search index {vs_index.display_name} at endpoint {vs_endpoint.display_name} ...\"\n",
        "        )\n",
        "        vs_deployed_index = vs_endpoint.deploy_index(\n",
        "            index=vs_index,\n",
        "            deployed_index_id=index_name,\n",
        "            display_name=index_name,\n",
        "            machine_type=\"e2-standard-16\",\n",
        "            min_replica_count=1,\n",
        "            max_replica_count=1,\n",
        "        )\n",
        "        print(\n",
        "            f\"Vector Search index {vs_index.display_name} is deployed at endpoint {vs_deployed_index.display_name}\"\n",
        "        )\n",
        "    else:\n",
        "        vs_deployed_index = aiplatform.MatchingEngineIndexEndpoint(\n",
        "            index_endpoint_name=index_endpoints[0][0]\n",
        "        )\n",
        "        print(\n",
        "            f\"Vector Search index {vs_index.display_name} is already deployed at endpoint {vs_deployed_index.display_name}\"\n",
        "        )\n",
        "\n",
        "    return vs_deployed_index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "hFGFmm4EmzlJ"
      },
      "outputs": [],
      "source": [
        "def setup():\n",
        "    # The number of dimensions for the gecko text embeddings is 768\n",
        "    VS_DIMENSIONS = 768\n",
        "    # Vertex AI Vector Search Index configuration\n",
        "\n",
        "    aiplatform.init(project=PROJECT_ID, location=REGION)\n",
        "\n",
        "    GCS_BUCKET_URI = f\"gs://{GCS_BUCKET}\"\n",
        "    new_bucket = create_bucket_class_location(GCS_BUCKET_URI)\n",
        "    vs_index = create_vector_search_index(VS_INDEX_NAME, VS_DIMENSIONS)\n",
        "    vs_endpoint = create_vector_search_endpoint(VS_INDEX_ENDPOINT_NAME)\n",
        "    vs_deployed_index = deploy_vector_search_endpoint(\n",
        "        vs_index, vs_endpoint, VS_INDEX_NAME\n",
        "    )\n",
        "\n",
        "    return new_bucket, vs_index, vs_endpoint, vs_deployed_index"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EUYC_nySmzlK"
      },
      "source": [
        "## Ingest data using LlamaIndex into Vertex AI Vector Search\n",
        "\n",
        "The following section leverages LlamaIndex to ingest, chunk, and embed the PDF data to be connected to the Vertex AI Vector Store.\n",
        "\n",
        "At the end of this section you will be ready to query against the Vector Store to find relevant context."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "jUK91ivpmzlK"
      },
      "outputs": [],
      "source": [
        "def initialize_llm_and_storage(\n",
        "    vs_index: aiplatform.MatchingEngineIndex,\n",
        "    vs_endpoint: aiplatform.MatchingEngineIndexEndpoint,\n",
        ") -> StorageContext:\n",
        "    \"\"\"Initializes Vertex AI Vector Store given a Vector Search index and deployed endpoint.\n",
        "    Configures embedding and LLMs models to be gecko and Gemini.\n",
        "    \"\"\"\n",
        "    # setup storage\n",
        "    vector_store = VertexAIVectorStore(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        index_id=vs_index.resource_name,\n",
        "        endpoint_id=vs_endpoint.resource_name,\n",
        "        gcs_bucket_name=GCS_BUCKET,\n",
        "    )\n",
        "\n",
        "    # set storage context\n",
        "    storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
        "\n",
        "    gemini_embedding_model = VertexTextEmbedding(\"text-embedding-005\")\n",
        "    llm = Vertex(\"gemini-2.0-flash\")\n",
        "\n",
        "    Settings.embed_model = gemini_embedding_model\n",
        "    Settings.llm = llm\n",
        "\n",
        "    return storage_context"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3g7boWv-y9BP"
      },
      "source": [
        "Setup Vector Search if it does not exist, otherwise return an endpoint for the deployed index, which will be used for the LlamaIndex retriever.\n",
        "\n",
        "Using SimpleDirectoryReade to read in the documents, which creates documents out of every file in a given directory. It is built into LlamaIndex and can read a variety of formats including Markdown, PDFs, Word documents, PowerPoint decks, images, audio and video."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9p4yqM99mzlK"
      },
      "outputs": [],
      "source": [
        "(bucket, vs_index, vs_endpoint, deployed_endpoint) = setup()\n",
        "storage_context = initialize_llm_and_storage(vs_index, vs_endpoint)\n",
        "docs = SimpleDirectoryReader(DOC_FOLDER).load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R07029k8mzlL"
      },
      "source": [
        "## Perform Q/A RAG\n",
        "\n",
        "This section performs the RAG prompt and returns an answer to the user query. To explore RAG frameworks we will look at 3 different options:\n",
        "\n",
        "1. Using the built-in RAG prompt provided by LlamaIndex\n",
        "2. Connecting a LangChain RAG template\n",
        "3. Creating a custom few-shot example RAG template.\n",
        "\n",
        "For each option, you will see the current text prompt structure and the generated output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "bHoRsvvzmzlL"
      },
      "outputs": [],
      "source": [
        "# Setting up helper functions\n",
        "\n",
        "\n",
        "def display_prompt_dict(prompts_dict: dict[str, BasePromptTemplate]):\n",
        "    \"\"\"Used to display the underlying text prompt used for RAG.\"\"\"\n",
        "    for k, p in prompts_dict.items():\n",
        "        text_md = f\"**Prompt Key**: {k}<br>**Text:** <br>\"\n",
        "        print(text_md)\n",
        "        print(p.get_template())\n",
        "        print(\"\\n\\n\")\n",
        "\n",
        "\n",
        "def display_and_run_prompt(query_engine: BaseQueryEngine, query_str: str):\n",
        "    \"\"\"Displays the current RAG prompt used and runs the query against the RAG workflow.\"\"\"\n",
        "    print(\"----Displaying current prompt dictionary----\\n\")\n",
        "    prompts_dict = query_engine.get_prompts()\n",
        "    display_prompt_dict(prompts_dict)\n",
        "\n",
        "    response = query_engine.query(query_str)\n",
        "    print(\"Response:\")\n",
        "    print(\"-\" * 80)\n",
        "    print(response.response)\n",
        "    print(\"-\" * 80)\n",
        "    print(\"Source Documents:\")\n",
        "    print(\"-\" * 80)\n",
        "    for source in response.source_nodes:\n",
        "        print(f\"Sample Text: {source.text[:200]}\")\n",
        "        print(f\"Relevance score: {source.get_score():.3f}\")\n",
        "        print(f\"File Name: {source.metadata.get('file_name')}\")\n",
        "        print(f\"Page #: {source.metadata.get('page_label')}\")\n",
        "        print(f\"File Path: {source.metadata.get('file_path')}\")\n",
        "        print(\"-\" * 80)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8bS5paNumzlL"
      },
      "source": [
        "### LlamaIndex Built-in RAG"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "d9KplPGrmzlL"
      },
      "outputs": [],
      "source": [
        "def llama_built_in_prompt(query_engine: BaseQueryEngine, query_str: str):\n",
        "    display_and_run_prompt(query_engine, query_str)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a9Sl1UBimzlL"
      },
      "source": [
        "### Templated RAG through LangChain"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "vx5metMvmzlL"
      },
      "outputs": [],
      "source": [
        "def langchain_rag_prompt(query_engine: BaseQueryEngine, query_str: str):\n",
        "    langchain_prompt = hub.pull(\"rlm/rag-prompt\")\n",
        "\n",
        "    langchain_prompt_template = LangchainPromptTemplate(\n",
        "        template=langchain_prompt,\n",
        "        template_var_mappings={\"query_str\": \"question\", \"context_str\": \"context\"},\n",
        "    )\n",
        "\n",
        "    query_engine.update_prompts(\n",
        "        {\"response_synthesizer:text_qa_template\": langchain_prompt_template}\n",
        "    )\n",
        "\n",
        "    display_and_run_prompt(query_engine, query_str)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eKCbqF85mzlL"
      },
      "source": [
        "### Custom RAG Implementation\n",
        "\n",
        "This custom RAG prompt highlights two important prompt engineering techniques:\n",
        "\n",
        "__Few-shot examples:__ Providing the model with a few examples of the desired input-output behavior helps guide the model's response, effectively demonstrating the task and the expected format. This is particularly useful when the task is complex or requires a specific style of output.\n",
        "\n",
        "__Grounding the output:__ Instructing the model to base its answer on the retrieved documents and to provide justification for the answer ensures that the response is factually grounded and relevant to the context. This is crucial for maintaining accuracy and preventing the model from generating responses that are either irrelevant or factually incorrect."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "jmjCfQhWmzlL"
      },
      "outputs": [],
      "source": [
        "def custom_few_shot_prompt(query_engine: BaseQueryEngine, query_str: str):\n",
        "    \"\"\"Generating custom few shot prompt to show the desired output format and prevent hallucination by including reasoning in the response.\"\"\"\n",
        "    qa_prompt_custom_string = \"\"\"\\\n",
        "    Context information is below.\n",
        "    ---------------------\n",
        "    {context_str}\n",
        "    ---------------------\n",
        "    Given the context information and not prior knowledge, answer the query asking about citations over different topics.\n",
        "\n",
        "    Please output your answer in the following JSON format:\n",
        "    JSON Output: [\n",
        "    \"answer\": 'This is the answer to the question',\n",
        "    \"justification\": 'This is the reasoning or evidence supporting the answer given the provided context'\n",
        "    ]\n",
        "\n",
        "    Example query and JSON output:\n",
        "    Query: Who are the authors of the paper?\n",
        "    JSON Output: [\n",
        "    \"answer\": \"The authors are Hikaru Fujita, Masaru Kashiwakura, Akihiro Kawagoe, Hisaki Hamamoto, Tetsuo Niitsuma, and Yuzuru Mitani.\"\n",
        "    \"justification\": \"The authors are listed on the first and last page in order.\"\n",
        "    ]\n",
        "\n",
        "    Query: When was there a failure at the Nakoso Power Station?\n",
        "    JSON Output: [\n",
        "    \"answer\": \"September  16,  2021\",\n",
        "    \"justification\": \"In the context provided it states: It was in this context\n",
        "    that  on  September  16,  2021,  a  failure  due  to  aging  forced  the  emergency  stop  of  the\n",
        "    unit No. 8 generator step-up transformer (built in 1981) at the Nakoso Power Station of\n",
        "    Jōban Joint Power Co., Ltd. \"\n",
        "    ]\n",
        "\n",
        "    Query: {query_str}\n",
        "    Answer:\n",
        "    \"\"\"\n",
        "\n",
        "    custom_RAG_template = PromptTemplate(template=qa_prompt_custom_string)\n",
        "\n",
        "    query_engine.update_prompts(\n",
        "        {\"response_synthesizer:text_qa_template\": custom_RAG_template}\n",
        "    )\n",
        "\n",
        "    display_and_run_prompt(query_engine, query_str)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "JwbEBAaKmzlM"
      },
      "outputs": [],
      "source": [
        "def index_and_query_documents(\n",
        "    documents: list[Document], storage_context: StorageContext\n",
        "):\n",
        "    \"\"\"Sets up vector store index to query against for a RAG pattern.\"\"\"\n",
        "    # Using Gemini embedding models\n",
        "    vector_index = VectorStoreIndex.from_documents(\n",
        "        documents, storage_context=storage_context\n",
        "    )\n",
        "\n",
        "    # Set up a query engine\n",
        "    query_engine = vector_index.as_query_engine()\n",
        "\n",
        "    query = \"what is minimum reserve rate of power?\"\n",
        "    print(\"*******Option 1: LlamaIndex Built-In Prompt*******\")\n",
        "    llama_built_in_prompt(query_engine, query)\n",
        "    print(\"*******Option 2: LangChain Template RAG Prompt*******\")\n",
        "    langchain_rag_prompt(query_engine, query)\n",
        "    print(\"*******Option 3: Custom Few-Shot Prompt*******\")\n",
        "    custom_few_shot_prompt(query_engine, query)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bk5CiUjumzlM"
      },
      "outputs": [],
      "source": [
        "# Run the RAG workflow for the LlamaIndex built in prompt, templated LangChain prompt, and custom few-shot prompt\n",
        "vector_idx = index_and_query_documents(docs, storage_context)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "96524db014e6"
      },
      "source": [
        "## Multi-Document RAG\n",
        "\n",
        "So far we have established how to set up set up the necessary Google Cloud resources, set up a LlamaIndex agent, and customize prompts for RAG question answering. \n",
        "\n",
        "In this section we will cover Multi-Document Agents that can effectively answer different set of questions over a larger set of documents. Questions can include QA and summaries over a individual document or across documents. \n",
        "\n",
        "To do this we will follow these steps:\n",
        "\n",
        "+ setup a \"document agent\" over each Document: each doc agent can do QA/summarization within its doc\n",
        "+ setup a top-level agent over this set of document agents: tool retrieval and answer over the set of tools responses to answer a question"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5abb721e2ecc"
      },
      "source": [
        "### Add more documents"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "eab28887f857"
      },
      "outputs": [],
      "source": [
        "def ingest_multi_document():\n",
        "    import os\n",
        "\n",
        "    doc_dict = {}\n",
        "\n",
        "    for filename in os.listdir(DOC_FOLDER):\n",
        "        doc_dict[filename] = SimpleDirectoryReader(\n",
        "            input_files=[os.path.join(DOC_FOLDER, filename)]\n",
        "        ).load_data()\n",
        "\n",
        "    return doc_dict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0db32726a93b"
      },
      "outputs": [],
      "source": [
        "# Load documents\n",
        "multi_docs = ingest_multi_document()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6a9394f7f0d3"
      },
      "source": [
        "### Build Document Level Agents\n",
        "\n",
        "First we will build document agents for each document. \n",
        "\n",
        "We will create two query engines, one for semantic search and one for summarization, for each document. These query engines will be converted into tools that can be passed to a function calling agent.\n",
        "\n",
        "We will be using the ReAct agent (short for \"Reasoning and Acting\"). This agent is an LLM-powered agent designed to perform complex tasks over your data. It operates in both \"read\" and \"write\" modes, making it a versatile tool for various applications. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 81,
      "metadata": {
        "id": "4a1958298fab"
      },
      "outputs": [],
      "source": [
        "def build_document_level_agents(documents, storage_context):\n",
        "    \"\"\"Sets up a vector search and summarization tool for each document. Generates an agent for each documents based on tools.\"\"\"\n",
        "    node_parser = SentenceSplitter()\n",
        "\n",
        "    # Build agents dictionary\n",
        "    agents = {}\n",
        "    query_engines = {}\n",
        "\n",
        "    for idx, doc_title in enumerate(documents):\n",
        "        # A Node represents a \"chunk\" of a source Document, whether that is a text chunk, an image, or other. Similar to Documents, they contain metadata and relationship information with other nodes.\n",
        "        nodes = node_parser.get_nodes_from_documents(\n",
        "            documents[doc_title], show_progress=True\n",
        "        )\n",
        "\n",
        "        # Build query index\n",
        "        vector_index = VectorStoreIndex.from_documents(\n",
        "            documents[doc_title], storage_context=storage_context\n",
        "        )\n",
        "        # Build summary index\n",
        "        summary_index = SummaryIndex(nodes)\n",
        "\n",
        "        # Define engines\n",
        "        vector_query_engine = vector_index.as_query_engine()\n",
        "        summary_query_engine = summary_index.as_query_engine()\n",
        "\n",
        "        # Define tools\n",
        "        query_engine_tools = [\n",
        "            QueryEngineTool(\n",
        "                query_engine=vector_query_engine,\n",
        "                metadata=ToolMetadata(\n",
        "                    name=\"vector_tool\",\n",
        "                    description=(\n",
        "                        f\"Useful for questions related to specific aspects of {doc_title}.\"\n",
        "                    ),\n",
        "                ),\n",
        "            ),\n",
        "            QueryEngineTool(\n",
        "                query_engine=summary_query_engine,\n",
        "                metadata=ToolMetadata(\n",
        "                    name=\"summary_tool\",\n",
        "                    description=(\n",
        "                        f\"Useful for any requests that require a holistic summary of EVERYTHING about {doc_title}. For questions about more specific sections, please use the vector_tool.\"\n",
        "                    ),\n",
        "                ),\n",
        "            ),\n",
        "        ]\n",
        "\n",
        "        # Build agent\n",
        "        llm = Vertex(\"gemini-2.0-flash\")\n",
        "        agent = ReActAgent.from_tools(\n",
        "            query_engine_tools,\n",
        "            llm=llm,\n",
        "            verbose=True,\n",
        "            system_prompt=f\"\"\"\\\n",
        "            You are a specialized agent designed to answer queries about {doc_title}.\n",
        "            You must ALWAYS use at least one of the tools provided when answering a question; do NOT rely on prior knowledge.\\\n",
        "            \"\"\",\n",
        "        )\n",
        "\n",
        "        agents[doc_title] = agent\n",
        "        query_engines[doc_title] = vector_index.as_query_engine(similarity_top_k=2)\n",
        "\n",
        "    return agents"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "43d555ed8d88"
      },
      "outputs": [],
      "source": [
        "agents = build_document_level_agents(multi_docs, storage_context)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5c870e09455e"
      },
      "source": [
        "### Build Top Level Agent\n",
        "\n",
        "We build a top-level agent that can orchestrate across the different document agents to answer any user query. This agent takes in all document agents as tools that were built above. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 84,
      "metadata": {
        "id": "8fb273be3ef1"
      },
      "outputs": [],
      "source": [
        "def build_top_level_agent(agents: dict[str, ReActAgent]) -> ReActAgent:\n",
        "    # This agent takes in all document agents as tools\n",
        "    all_tools = []\n",
        "\n",
        "    for filename in os.listdir(DOC_FOLDER):\n",
        "        summary = f\"This content contains a research paper articles about {filename}. Use this tool if you want to answer any questions about {filename}.\\n\"\n",
        "        doc_tool = QueryEngineTool(\n",
        "            query_engine=agents[filename],\n",
        "            metadata=ToolMetadata(\n",
        "                name=f\"tool_{filename}\".rstrip(\".pdf\"),\n",
        "                description=summary,\n",
        "            ),\n",
        "        )\n",
        "\n",
        "        all_tools.append(doc_tool)\n",
        "\n",
        "    # define an \"object\" index and retriever over these tools\n",
        "    obj_index = ObjectIndex.from_objects(\n",
        "        all_tools,\n",
        "        index_cls=VectorStoreIndex,\n",
        "    )\n",
        "\n",
        "    # Create top level agent\n",
        "    top_agent = ReActAgent.from_tools(\n",
        "        tool_retriever=obj_index.as_retriever(similarity_top_k=3),\n",
        "        system_prompt=\"\"\" \\\n",
        "            You are an agent designed to answer queries about energy systems.\n",
        "            Please always use the tools provided to answer a question. Do not rely on prior knowledge.\\\n",
        "\n",
        "            \"\"\",\n",
        "        verbose=True,\n",
        "    )\n",
        "\n",
        "    return top_agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 85,
      "metadata": {
        "id": "6015192f0c2c"
      },
      "outputs": [],
      "source": [
        "top_level_agent = build_top_level_agent(agents)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "df1cd003853d"
      },
      "source": [
        "### Perform Multi-Document QA Rag\n",
        "\n",
        "Know we can query the top level agent. We will experiment with various question types that require lookup in individual documents and across multiple documents."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 86,
      "metadata": {
        "id": "9e5df3e2612f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
            "Action: tool_04a04\n",
            "Action Input: {'input': 'Higashi-Shimizu Substation'}\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
            "Action: vector_tool\n",
            "Action Input: {'input': 'Higashi-Shimizu Substation'}\n",
            "\u001b[0m\u001b[1;3;34mObservation: ## Higashi-Shimizu Substation: A Compact Powerhouse\n",
            "\n",
            "The Higashi-Shimizu Substation is a testament to efficient engineering and collaboration. Located on a steep hillside with limited space, the substation houses three frequency converters (FCs): No. 1 FC, No. 2 FC, and No. 3 FC. \n",
            "\n",
            "**Key Features:**\n",
            "\n",
            "* **Compact Design:** Despite the limited space, the substation boasts a compact design that optimizes equipment layout while maintaining required electrical and maintenance clearances.\n",
            "* **Redundant Control and Protection Systems:** The MACH system ensures reliable operation with redundant control and protection computers. Standby redundancy for control and parallel redundancy for protection guarantee uninterrupted operation even during maintenance.\n",
            "* **Advanced Frequency Converters:** The substation utilizes advanced frequency converters with impressive specifications:\n",
            "    * Interconnection point voltage: 200 kV (phase voltage)\n",
            "    * Interconnection point current: 2000 kA\n",
            "    * Active power: 0 MW\n",
            "    * Reactive power: 0.2 Mvar\n",
            "    * Active current: 0.51 p.u.\n",
            "    * Reactive current: 0.51 p.u.\n",
            "* **Fault Tolerance:** The converters can continue operation without gate blocking even during malfunctions like short circuits or grounding faults on the AC grid side. Additionally, they contribute to power grid stabilization by injecting reactive current during grid failures.\n",
            "* **Building Design:** The FC building will have one underground floor and two floors above ground, with a total height of approximately 25 meters.\n",
            "\n",
            "**Future Outlook:**\n",
            "\n",
            "Hitachi positions HVDC as a core business within the energy field and aims to meet the growing demand for grid interconnection facilities. The Higashi-Shimizu Substation exemplifies Hitachi's commitment to contributing to a decarbonized society by facilitating the integration of renewable energy sources and enhancing grid stability. \n",
            "\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
            "Answer: ## Higashi-Shimizu Substation: A Compact Powerhouse\n",
            "\n",
            "The Higashi-Shimizu Substation is a testament to efficient engineering and collaboration. Located on a steep hillside with limited space, the substation houses three frequency converters (FCs): No. 1 FC, No. 2 FC, and No. 3 FC. \n",
            "\n",
            "**Key Features:**\n",
            "\n",
            "* **Compact Design:** Despite the limited space, the substation boasts a compact design that optimizes equipment layout while maintaining required electrical and maintenance clearances.\n",
            "* **Redundant Control and Protection Systems:** The MACH system ensures reliable operation with redundant control and protection computers. Standby redundancy for control and parallel redundancy for protection guarantee uninterrupted operation even during maintenance.\n",
            "* **Advanced Frequency Converters:** The substation utilizes advanced frequency converters with impressive specifications:\n",
            "    * Interconnection point voltage: 200 kV (phase voltage)\n",
            "    * Interconnection point current: 2000 kA\n",
            "    * Active power: 0 MW\n",
            "    * Reactive power: 0.2 Mvar\n",
            "    * Active current: 0.51 p.u.\n",
            "    * Reactive current: 0.51 p.u.\n",
            "* **Fault Tolerance:** The converters can continue operation without gate blocking even during malfunctions like short circuits or grounding faults on the AC grid side. Additionally, they contribute to power grid stabilization by injecting reactive current during grid failures.\n",
            "* **Building Design:** The FC building will have one underground floor and two floors above ground, with a total height of approximately 25 meters.\n",
            "\n",
            "**Future Outlook:**\n",
            "\n",
            "Hitachi positions HVDC as a core business within the energy field and aims to meet the growing demand for grid interconnection facilities. The Higashi-Shimizu Substation exemplifies Hitachi's commitment to contributing to a decarbonized society by facilitating the integration of renewable energy sources and enhancing grid stability.\n",
            "\u001b[0m\u001b[1;3;34mObservation: ## Higashi-Shimizu Substation: A Compact Powerhouse\n",
            "\n",
            "The Higashi-Shimizu Substation is a testament to efficient engineering and collaboration. Located on a steep hillside with limited space, the substation houses three frequency converters (FCs): No. 1 FC, No. 2 FC, and No. 3 FC. \n",
            "\n",
            "**Key Features:**\n",
            "\n",
            "* **Compact Design:** Despite the limited space, the substation boasts a compact design that optimizes equipment layout while maintaining required electrical and maintenance clearances.\n",
            "* **Redundant Control and Protection Systems:** The MACH system ensures reliable operation with redundant control and protection computers. Standby redundancy for control and parallel redundancy for protection guarantee uninterrupted operation even during maintenance.\n",
            "* **Advanced Frequency Converters:** The substation utilizes advanced frequency converters with impressive specifications:\n",
            "    * Interconnection point voltage: 200 kV (phase voltage)\n",
            "    * Interconnection point current: 2000 kA\n",
            "    * Active power: 0 MW\n",
            "    * Reactive power: 0.2 Mvar\n",
            "    * Active current: 0.51 p.u.\n",
            "    * Reactive current: 0.51 p.u.\n",
            "* **Fault Tolerance:** The converters can continue operation without gate blocking even during malfunctions like short circuits or grounding faults on the AC grid side. Additionally, they contribute to power grid stabilization by injecting reactive current during grid failures.\n",
            "* **Building Design:** The FC building will have one underground floor and two floors above ground, with a total height of approximately 25 meters.\n",
            "\n",
            "**Future Outlook:**\n",
            "\n",
            "Hitachi positions HVDC as a core business within the energy field and aims to meet the growing demand for grid interconnection facilities. The Higashi-Shimizu Substation exemplifies Hitachi's commitment to contributing to a decarbonized society by facilitating the integration of renewable energy sources and enhancing grid stability.\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
            "Answer: The Higashi-Shimizu Substation is a compact and efficient substation located on a steep hillside with limited space. It houses three frequency converters (FCs) and features a redundant control and protection system, advanced frequency converters, and fault tolerance. The substation is designed to contribute to a decarbonized society by facilitating the integration of renewable energy sources and enhancing grid stability.\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "# QA over a specific doc\n",
        "response = top_level_agent.query(\n",
        "    \"Tell me about the Higashi-Shimizu Substation of Chubu Electric Power Grid\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 89,
      "metadata": {
        "id": "dd041fd45c18"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
            "Action: tool_04a01\n",
            "Action Input: {'input': 'What are all projects introduced after the Great East Japan Earthquake?'}\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
            "Action: summary_tool\n",
            "Action Input: {'input': 'What are all projects introduced after the Great East Japan Earthquake?'}\n",
            "\u001b[0m\u001b[1;3;34mObservation: ## Projects introduced after the Great East Japan Earthquake:\n",
            "\n",
            "* **Enhancement of Higashi-Shimizu Substation of Chubu Electric Power Grid by VSC Technology:** This project aims to expand the grid interconnection capacity between the different frequencies in eastern Japan (50 Hz) and western Japan (60 Hz) by installing two blocks of frequency converter systems at the Higashi-Shimizu Substation. The project is expected to be completed by the end of FY2027.\n",
            "* **Hida-Shinano HVDC link:** This project increased the interconnection capacity between the 60 Hz Chubu Electric Power Grid area and the 50 Hz area from the initial 1.2 million kW to 2.1 million kW.\n",
            "* **Construction of a new 275-kV power transmission line on the 50 Hz side of the Higashi-Shimizu Substation:** This project will connect the newly-installed No. 1 FC and No. 3 FC to the existing No. 2 FC.\n",
            "* **Installation of a new air-insulated filter for HVDC systems at the Hida Converter Station of Chubu Electric Power Grid Co., Inc.:** This project was the first overseas-made air-insulated filter for HVDC systems delivered by Hitachi.\n",
            "\n",
            "These projects are all aimed at improving the resilience of the power grid and enabling the expansion of renewable energy sources.\n",
            "\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
            "Answer: The projects introduced after the Great East Japan Earthquake are:\n",
            "\n",
            "* Enhancement of Higashi-Shimizu Substation of Chubu Electric Power Grid by VSC Technology\n",
            "* Hida-Shinano HVDC link\n",
            "* Construction of a new 275-kV power transmission line on the 50 Hz side of the Higashi-Shimizu Substation\n",
            "* Installation of a new air-insulated filter for HVDC systems at the Hida Converter Station of Chubu Electric Power Grid Co., Inc.\n",
            "\n",
            "These projects are all aimed at improving the resilience of the power grid and enabling the expansion of renewable energy sources.\n",
            "\u001b[0m\u001b[1;3;34mObservation: The projects introduced after the Great East Japan Earthquake are:\n",
            "\n",
            "* Enhancement of Higashi-Shimizu Substation of Chubu Electric Power Grid by VSC Technology\n",
            "* Hida-Shinano HVDC link\n",
            "* Construction of a new 275-kV power transmission line on the 50 Hz side of the Higashi-Shimizu Substation\n",
            "* Installation of a new air-insulated filter for HVDC systems at the Hida Converter Station of Chubu Electric Power Grid Co., Inc.\n",
            "\n",
            "These projects are all aimed at improving the resilience of the power grid and enabling the expansion of renewable energy sources.\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
            "Answer: The projects introduced after the Great East Japan Earthquake are:\n",
            "\n",
            "* Enhancement of Higashi-Shimizu Substation of Chubu Electric Power Grid by VSC Technology\n",
            "* Hida-Shinano HVDC link\n",
            "* Construction of a new 275-kV power transmission line on the 50 Hz side of the Higashi-Shimizu Substation\n",
            "* Installation of a new air-insulated filter for HVDC systems at the Hida Converter Station of Chubu Electric Power Grid Co., Inc.\n",
            "\n",
            "These projects are all aimed at improving the resilience of the power grid and enabling the expansion of renewable energy sources.\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "# summaries across documents\n",
        "response = top_level_agent.query(\n",
        "    \"What are all projects introduced after the Great East Japan Earthquake?\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 91,
      "metadata": {
        "id": "535d5c650351"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
            "Action: tool_04a04\n",
            "Action Input: {'input': 'List out all the technologies that are used to stabilize the power system?'}\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
            "Action: summary_tool\n",
            "Action Input: {'input': 'List out all the technologies that are used to stabilize the power system?'}\n",
            "\u001b[0m\u001b[1;3;34mObservation: ## Technologies used to stabilize the power system:\n",
            "\n",
            "Based on the provided context, here are the technologies used to stabilize the power system:\n",
            "\n",
            "* **HVDC (High-Voltage Direct Current) technology:** This technology allows for the transmission of large amounts of electricity over long distances with minimal losses. It also enables the integration of renewable energy sources and the interconnection of grids, which helps to stabilize the overall power system.\n",
            "* **HVDC Classic and HVDC Light:** These are two specific types of HVDC technology developed by Hitachi Energy. They offer high efficiency, controllability, and the ability to solve many of the challenges in today's complex power grids.\n",
            "* **MACH control system:** This is a state-of-the-art control system developed by Hitachi Energy that provides unequalled calculation capacity and enables a high degree of integration and handling of all HVDC control and protection functions. It is key to the superior performance and reliability of HVDC and power quality solutions supplied by Hitachi Energy.\n",
            "* **Converter control:** HVDC converters can be controlled to provide fast power changes that can be used to stabilize connected AC power systems in many different ways. This includes damping power oscillations, providing support for sudden loss of power generation, and balancing voltage changes using reactive power control.\n",
            "* **Fast-reacting protections:** HVDC systems require protections that are much faster than traditional AC protections to avoid damage to the converter valves. The MACH system includes a fully duplicated centralized protection system that meets these requirements.\n",
            "* **Transient fault recorders:** These devices are integrated into the MACH system to generate easily readable records of faults in the system. This information can be used to improve the reliability of the power system.\n",
            "* **Line and cable fault locators:** These devices use travelling waves technology to accurately determine the location of faults in DC lines or cables. This information can be used to quickly restore power after a fault occurs.\n",
            "\n",
            "In addition to these technologies, there are many other technologies that can be used to stabilize the power system. These include:\n",
            "\n",
            "* **Energy storage:** Energy storage can be used to store excess energy from renewable sources and release it when needed to help balance the grid.\n",
            "* **Demand response:** Demand response programs encourage consumers to reduce their electricity consumption during peak times, which can help to reduce the strain on the power system.\n",
            "* **Smart grids:** Smart grids use advanced technologies to monitor and control the power system in real-time, which can help to improve efficiency and reliability.\n",
            "\n",
            "By using a combination of these technologies\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
            "Answer: The following technologies are used to stabilize the power system:\n",
            "\n",
            "* **HVDC (High-Voltage Direct Current) technology:** This technology allows for the transmission of large amounts of electricity over long distances with minimal losses. It also enables the integration of renewable energy sources and the interconnection of grids, which helps to stabilize the overall power system.\n",
            "* **HVDC Classic and HVDC Light:** These are two specific types of HVDC technology developed by Hitachi Energy. They offer high efficiency, controllability, and the ability to solve many of the challenges in today's complex power grids.\n",
            "* **MACH control system:** This is a state-of-the-art control system developed by Hitachi Energy that provides unequalled calculation capacity and enables a high degree of integration and handling of all HVDC control and protection functions. It is key to the superior performance and reliability of HVDC and power quality solutions supplied by Hitachi Energy.\n",
            "* **Converter control:** HVDC converters can be controlled to provide fast power changes that can be used to stabilize connected AC power systems in many different ways. This includes damping power oscillations, providing support for sudden loss of power generation, and balancing voltage changes using reactive power control.\n",
            "* **Fast-reacting protections:** HVDC systems require protections that are much faster than traditional AC protections to avoid damage to the converter valves. The MACH system includes a fully duplicated centralized protection system that meets these requirements.\n",
            "* **Transient fault recorders:** These devices are integrated into the MACH system to generate easily readable records of faults in the system. This information can be used to improve the reliability of the power system.\n",
            "* **Line and cable fault locators:** These devices use travelling waves technology to accurately determine the location of faults in DC lines or cables. This information can be used to quickly restore power after a fault occurs.\n",
            "\n",
            "In addition to these technologies, there are many other technologies that can be used to stabilize the power system. These include:\n",
            "\n",
            "* **Energy storage:** Energy storage can be used to store excess energy from renewable sources and release it when needed to help balance the grid.\n",
            "* **Demand response:** Demand response programs encourage consumers to reduce their electricity consumption during peak times, which can help to reduce the strain on the power system.\n",
            "* **Smart grids:** Smart grids use advanced technologies to monitor and control the power system in real-time, which can help to improve efficiency and reliability\n",
            "\u001b[0m\u001b[1;3;34mObservation: The following technologies are used to stabilize the power system:\n",
            "\n",
            "* **HVDC (High-Voltage Direct Current) technology:** This technology allows for the transmission of large amounts of electricity over long distances with minimal losses. It also enables the integration of renewable energy sources and the interconnection of grids, which helps to stabilize the overall power system.\n",
            "* **HVDC Classic and HVDC Light:** These are two specific types of HVDC technology developed by Hitachi Energy. They offer high efficiency, controllability, and the ability to solve many of the challenges in today's complex power grids.\n",
            "* **MACH control system:** This is a state-of-the-art control system developed by Hitachi Energy that provides unequalled calculation capacity and enables a high degree of integration and handling of all HVDC control and protection functions. It is key to the superior performance and reliability of HVDC and power quality solutions supplied by Hitachi Energy.\n",
            "* **Converter control:** HVDC converters can be controlled to provide fast power changes that can be used to stabilize connected AC power systems in many different ways. This includes damping power oscillations, providing support for sudden loss of power generation, and balancing voltage changes using reactive power control.\n",
            "* **Fast-reacting protections:** HVDC systems require protections that are much faster than traditional AC protections to avoid damage to the converter valves. The MACH system includes a fully duplicated centralized protection system that meets these requirements.\n",
            "* **Transient fault recorders:** These devices are integrated into the MACH system to generate easily readable records of faults in the system. This information can be used to improve the reliability of the power system.\n",
            "* **Line and cable fault locators:** These devices use travelling waves technology to accurately determine the location of faults in DC lines or cables. This information can be used to quickly restore power after a fault occurs.\n",
            "\n",
            "In addition to these technologies, there are many other technologies that can be used to stabilize the power system. These include:\n",
            "\n",
            "* **Energy storage:** Energy storage can be used to store excess energy from renewable sources and release it when needed to help balance the grid.\n",
            "* **Demand response:** Demand response programs encourage consumers to reduce their electricity consumption during peak times, which can help to reduce the strain on the power system.\n",
            "* **Smart grids:** Smart grids use advanced technologies to monitor and control the power system in real-time, which can help to improve efficiency and reliability\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
            "Answer: The following technologies are used to stabilize the power system:\n",
            "\n",
            "* **HVDC (High-Voltage Direct Current) technology:** This technology allows for the transmission of large amounts of electricity over long distances with minimal losses. It also enables the integration of renewable energy sources and the interconnection of grids, which helps to stabilize the overall power system.\n",
            "* **HVDC Classic and HVDC Light:** These are two specific types of HVDC technology developed by Hitachi Energy. They offer high efficiency, controllability, and the ability to solve many of the challenges in today's complex power grids.\n",
            "* **MACH control system:** This is a state-of-the-art control system developed by Hitachi Energy that provides unequalled calculation capacity and enables a high degree of integration and handling of all HVDC control and protection functions. It is key to the superior performance and reliability of HVDC and power quality solutions supplied by Hitachi Energy.\n",
            "* **Converter control:** HVDC converters can be controlled to provide fast power changes that can be used to stabilize connected AC power systems in many different ways. This includes damping power oscillations, providing support for sudden loss of power generation, and balancing voltage changes using reactive power control.\n",
            "* **Fast-reacting protections:** HVDC systems require protections that are much faster than traditional AC protections to avoid damage to the converter valves. The MACH system includes a fully duplicated centralized protection system that meets these requirements.\n",
            "* **Transient fault recorders:** These devices are integrated into the MACH system to generate easily readable records of faults in the system. This information can be used to improve the reliability of the power system.\n",
            "* **Line and cable fault locators:** These devices use travelling waves technology to accurately determine the location of faults in DC lines or cables. This information can be used to quickly restore power after a fault occurs.\n",
            "\n",
            "In addition to these technologies, there are many other technologies that can be used to stabilize the power system. These include:\n",
            "\n",
            "* **Energy storage:** Energy storage can be used to store excess energy from renewable sources and release it when needed to help balance the grid.\n",
            "* **Demand response:** Demand response programs encourage consumers to reduce their electricity consumption during peak times, which can help to reduce the strain on the power system.\n",
            "* **Smart grids:** Smart grids use advanced technologies to monitor and control the power system in real-time, which can help to improve efficiency and reliability\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "# cross document QA\n",
        "response = top_level_agent.query(\n",
        "    \"List out all the technologies that are used to stabilize the power system?\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 94,
      "metadata": {
        "id": "3290504b6bc0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
            "Action: tool_04a04\n",
            "Action Input: {'input': 'Explain to me what the building blocks of the MACH control and protection system and where it is used.'}\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
            "Action: vector_tool\n",
            "Action Input: {'input': 'Explain to me what the building blocks of the MACH control and protection system and where it is used.'}\n",
            "\u001b[0m\u001b[1;3;34mObservation: The MACH control and protection system is a modular system that can control all types and generations of HVDC systems. It is used to monitor, control, and protect HVDC systems. The system is made up of several different building blocks, including:\n",
            "\n",
            "* The main computer system, which consists of three different products: the PS700, a high-performance server type computer with 8–12 cores; the PS935, an 8-core DSP platform; and the PS775 redundancy changeover unit.\n",
            "* The input/output (I/O) system, which receives input from the main computer system and sends output to the converter valve electronics.\n",
            "* The converter valve electronics, which connect directly to the semiconductors in the converter valve.\n",
            "* The valve control unit, which acts as an interface between the converter valve electronics and the main control system.\n",
            "\n",
            "These building blocks work together to form control and protection loops that range from tens of microseconds to milliseconds. The system is programmed using a graphical programming language called HiDraw, which allows for design, real-time debugging, and simulations of complete converter systems before they are built.\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
            "Answer: The MACH control and protection system is a modular system that can control all types and generations of HVDC systems. It is used to monitor, control, and protect HVDC systems. The system is made up of several different building blocks, including the main computer system, the input/output (I/O) system, the converter valve electronics, and the valve control unit. These building blocks work together to form control and protection loops that range from tens of microseconds to milliseconds. The system is programmed using a graphical programming language called HiDraw, which allows for design, real-time debugging, and simulations of complete converter systems before they are built. The MACH control and protection system is used in a variety of applications, including power transmission, wind power, and offshore oil and gas platforms.\n",
            "\u001b[0m\u001b[1;3;34mObservation: The MACH control and protection system is a modular system that can control all types and generations of HVDC systems. It is used to monitor, control, and protect HVDC systems. The system is made up of several different building blocks, including the main computer system, the input/output (I/O) system, the converter valve electronics, and the valve control unit. These building blocks work together to form control and protection loops that range from tens of microseconds to milliseconds. The system is programmed using a graphical programming language called HiDraw, which allows for design, real-time debugging, and simulations of complete converter systems before they are built. The MACH control and protection system is used in a variety of applications, including power transmission, wind power, and offshore oil and gas platforms.\n",
            "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
            "Answer: The MACH control and protection system is a modular system that can control all types and generations of HVDC systems. It is used to monitor, control, and protect HVDC systems. The system is made up of several different building blocks, including the main computer system, the input/output (I/O) system, the converter valve electronics, and the valve control unit. These building blocks work together to form control and protection loops that range from tens of microseconds to milliseconds. The system is programmed using a graphical programming language called HiDraw, which allows for design, real-time debugging, and simulations of complete converter systems before they are built. The MACH control and protection system is used in a variety of applications, including power transmission, wind power, and offshore oil and gas platforms.\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "response = top_level_agent.query(\n",
        "    \"Explain to me what the building blocks of the MACH control and protection system and where it is used.\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lXlQhBl_mzlM"
      },
      "source": [
        "## Cleaning up\n",
        "To clean up all Google Cloud resources used in this project, you can delete the Google Cloud project you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "llamaindex_rag.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
