gemini/orchestration/llamaindex_workflows.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "09472b4cfe54"
      },
      "source": [
        "# LlamaIndex RAG Workflows using Gemini and Firestore\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/orchestration/llamaindex_workflows.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%2Forchestration%2Fllamaindex_workflows.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/orchestration/llamaindex_workflows.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/orchestration/llamaindex_workflows.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/orchestration/llamaindex_workflows.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/orchestration/llamaindex_workflows.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/orchestration/llamaindex_workflows.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/orchestration/llamaindex_workflows.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/orchestration/llamaindex_workflows.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": "18b1887a153f"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) | [Noa Ben-Efraim](https://github.com/noabenefraim) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4dddad16697f"
      },
      "source": [
        "## Overview\n",
        "LlamaIndex workflows are a powerful way to orchestrate complex LLM (large language model) applications. They provide an event-driven framework for building AI systems that go beyond simple question-answering.   \n",
        "\n",
        "Think of a workflow as a series of steps, where each step performs a specific action. These actions can be anything from querying an LLM, to retrieving data from a vector database, to interacting with external APIs. The workflow manages the flow of data between these steps, making it easy to build sophisticated AI applications.   \n",
        "\n",
        "Here's a breakdown of the key concepts:\n",
        "\n",
        "+ Events: These trigger actions within the workflow. For example, a user's query can be an initial event that kicks off the workflow.   \n",
        "+ Steps: These are individual functions decorated with @step that process events and potentially emit new events. Steps are the building blocks of your workflow.   \n",
        "+ Event-driven: This means that the workflow reacts to events as they happen, making it flexible and dynamic."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5bdceacc7b39"
      },
      "source": [
        "This notebook perform a complex Retrieval Augmented Generation (RAG) workflow using Gemini models and Firestore databases. There are two branches for this workflow:\n",
        "\n",
        "_Branch 1_\n",
        "+ Start Event triggered by providing a data directory to the workflow\n",
        "+ Ingest data using the LlamaIndex `SimpleDirectoryReader`\n",
        "+ Load data in the Firestore Database\n",
        "\n",
        "_Branch 2_\n",
        "+ Start Event triggered by providing a query to the workflow\n",
        "+ The QueryMultiStep Event that breaks down a complex query into sequential sub-questions using Gemini. Then proceeds to answer the sub-questions.\n",
        "+ The sub-questions results are passed to the RerankEvent where given the initial user query, Gemini reranks the returned answers to the sub-questions.\n",
        "+ The reranked chunks are passed to the CreateCitationEvents where citations are added to the sub-questions used to generate the answer.\n",
        "+ An answer is synthesized for the original query and returned to the user.\n",
        "\n",
        "References:\n",
        "+ https://docs.llamaindex.ai/en/stable/examples/workflow/rag/\n",
        "+ https://docs.llamaindex.ai/en/stable/examples/workflow/multi_step_query_engine/\n",
        "+ https://docs.llamaindex.ai/en/stable/examples/workflow/citation_query_engine/\n",
        "\n",
        "\n",
        "![RAGWorkflow](https://storage.googleapis.com/github-repo/generative-ai/gemini/orchestration/llamaindex_workflows/RAGWorkflow.png)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef012cf7cb67"
      },
      "source": [
        "## Get started\n",
        "\n",
        "### Install required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "323a27d12c02"
      },
      "outputs": [],
      "source": [
        "%pip install llama-index==\"0.11.8\" \\\n",
        "    llama-index-embeddings-vertex==\"0.2.0\" \\\n",
        "    llama-index-utils-workflow==\"0.2.1\" \\\n",
        "    llama-index-llms-vertex==\"0.3.4\" \\\n",
        "    llama-index-storage-docstore-firestore==\"0.2.0\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f42d12d15616"
      },
      "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": "f3d98395d9a4"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e114f5653870"
      },
      "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": "911453311a5d"
      },
      "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": "e8d7771a5818"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "This notebook requires the following resources:\n",
        "+ Initialized Google Cloud project\n",
        "+ Vertex AI API enabled\n",
        "+ Existing VPC/Subnet\n",
        "+ Existing Firestore database\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",
        "To get started using Firestore Database, refer to the following [documentation](https://cloud.google.com/firestore/docs/manage-databases).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "e04ae6146ccd"
      },
      "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",
        "FIRESTORE_DATABASE_ID = \"[your-firestore-database-id]\"\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d68eca7a8d4f"
      },
      "source": [
        "## Workflow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "93704f34a080"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "9b7824a768bc"
      },
      "outputs": [],
      "source": [
        "from typing import Any, cast\n",
        "\n",
        "from IPython.display import Markdown, display\n",
        "from llama_index.core import (\n",
        "    Settings,\n",
        "    SimpleDirectoryReader,\n",
        "    StorageContext,\n",
        "    VectorStoreIndex,\n",
        ")\n",
        "from llama_index.core.indices.query.query_transform.base import (\n",
        "    StepDecomposeQueryTransform,\n",
        ")\n",
        "from llama_index.core.llms import LLM\n",
        "from llama_index.core.node_parser import SentenceSplitter\n",
        "from llama_index.core.postprocessor.llm_rerank import LLMRerank\n",
        "from llama_index.core.prompts import PromptTemplate\n",
        "from llama_index.core.response_synthesizers import (\n",
        "    ResponseMode,\n",
        "    get_response_synthesizer,\n",
        ")\n",
        "from llama_index.core.schema import MetadataMode, NodeWithScore, QueryBundle, TextNode\n",
        "from llama_index.core.workflow import (\n",
        "    Context,\n",
        "    Event,\n",
        "    StartEvent,\n",
        "    StopEvent,\n",
        "    Workflow,\n",
        "    step,\n",
        ")\n",
        "from llama_index.embeddings.vertex import VertexTextEmbedding\n",
        "from llama_index.llms.vertex import Vertex\n",
        "from llama_index.storage.docstore.firestore import FirestoreDocumentStore\n",
        "from llama_index.utils.workflow import draw_all_possible_flows\n",
        "from vertexai.generative_models import HarmBlockThreshold, HarmCategory, SafetySetting"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0e101376082c"
      },
      "source": [
        "### Get data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f15a704ea11b"
      },
      "outputs": [],
      "source": [
        "!mkdir -p './data'\n",
        "!wget 'https://www.gutenberg.org/cache/epub/64317/pg64317.txt' -O 'data/gatsby.txt'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "927a7ffd9ad8"
      },
      "source": [
        "### Set credentials"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "7d11aac2947b"
      },
      "outputs": [],
      "source": [
        "import google.auth\n",
        "import google.auth.transport.requests\n",
        "\n",
        "# credentials will now have an api token\n",
        "credentials = google.auth.default(quota_project_id=PROJECT_ID)[0]\n",
        "request = google.auth.transport.requests.Request()\n",
        "credentials.refresh(request)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af476af08250"
      },
      "source": [
        "## Workflow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "456efd13df2a"
      },
      "source": [
        "### Set up the LLM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "695792a24ba9"
      },
      "outputs": [],
      "source": [
        "safety_config = [\n",
        "    SafetySetting(\n",
        "        category=HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,\n",
        "        threshold=HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    ),\n",
        "    SafetySetting(\n",
        "        category=HarmCategory.HARM_CATEGORY_HARASSMENT,\n",
        "        threshold=HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    ),\n",
        "    SafetySetting(\n",
        "        category=HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,\n",
        "        threshold=HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    ),\n",
        "]\n",
        "embedding_model = VertexTextEmbedding(\n",
        "    model_name=\"text-embedding-004\", credentials=credentials\n",
        ")\n",
        "llm = Vertex(\n",
        "    model=\"gemini-pro\",\n",
        "    temperature=0.2,\n",
        "    max_tokens=3000,\n",
        "    safety_settings=safety_config,\n",
        "    credentials=credentials,\n",
        ")\n",
        "\n",
        "Settings.embed_model = embedding_model\n",
        "Settings.llm = llm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1b99e7e4aa0d"
      },
      "source": [
        "### Define Event classes\n",
        "\n",
        "Here we will create custom events that can be emitted by steps and trigger other steps. \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "31173e6befe2"
      },
      "outputs": [],
      "source": [
        "class RetrieverEvent(Event):\n",
        "    \"\"\"Result of running retrieval\"\"\"\n",
        "\n",
        "    nodes: list[NodeWithScore]\n",
        "\n",
        "\n",
        "class RerankEvent(Event):\n",
        "    \"\"\"Result of running reranking on retrieved nodes\"\"\"\n",
        "\n",
        "    nodes: list[NodeWithScore]\n",
        "    source_nodes: list[NodeWithScore]\n",
        "    final_response_metadata: dict[str, Any]\n",
        "\n",
        "\n",
        "class FirestoreIndexData(Event):\n",
        "    \"\"\"Result of indexing documents in Firestore\"\"\"\n",
        "\n",
        "    status: str\n",
        "\n",
        "\n",
        "class QueryMultiStepEvent(Event):\n",
        "    \"\"\"\n",
        "    Event containing results of a multi-step query process.\n",
        "\n",
        "    Attributes:\n",
        "        nodes (List[NodeWithScore]): List of nodes with their associated scores.\n",
        "        source_nodes (List[NodeWithScore]): List of source nodes with their scores.\n",
        "        final_response_metadata (Dict[str, Any]): Metadata associated with the final response.\n",
        "    \"\"\"\n",
        "\n",
        "    nodes: list[NodeWithScore]\n",
        "    source_nodes: list[NodeWithScore]\n",
        "    final_response_metadata: dict[str, Any]\n",
        "\n",
        "\n",
        "class CreateCitationsEvent(Event):\n",
        "    \"\"\"Add citations to the nodes.\"\"\"\n",
        "\n",
        "    nodes: list[NodeWithScore]\n",
        "    source_nodes: list[NodeWithScore]\n",
        "    final_response_metadata: dict[str, Any]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "805eae826756"
      },
      "source": [
        "### Update Prompt Templates\n",
        "\n",
        "Defining custom prompts used for the citation portion of the workflow."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "b2ec8b2859f2"
      },
      "outputs": [],
      "source": [
        "CITATION_QA_TEMPLATE = PromptTemplate(\n",
        "    \"Your task is to answer the question based on the information given in the sources listed below.\"\n",
        "    \"Use only the provided sources to answer.\"\n",
        "    \"Cite the source number(s) for any information you use in your answer (e.g., [1]).\"\n",
        "    \"Always include at least one source citation in your answer.\"\n",
        "    \"Only cite a source if you directly use information from it.\"\n",
        "    \"If the sources don't contain the information needed to answer the question, state that.\"\n",
        "    \"For example:\"\n",
        "    \"Source 1: Apples are red, green, or yellow.\"\n",
        "    \"Source 2:  Bananas are yellow when ripe.\"\n",
        "    \"Source 3: Strawberries are red when ripe.\"\n",
        "    \"Query: Which fruits are red when ripe?\"\n",
        "    \"Answer: Apples [1] and strawberries [3] can be red when ripe.\"\n",
        "    \"------\"\n",
        "    \"Below are several numbered sources of information:\"\n",
        "    \"------\"\n",
        "    \"{context_str}\"\n",
        "    \"------\"\n",
        "    \"Query: {query_str}\"\n",
        "    \"Answer: \"\n",
        ")\n",
        "\n",
        "CITATION_REFINE_TEMPLATE = PromptTemplate(\n",
        "    \"You have an initial answer to a query.\"\n",
        "    \"Your job is to improve this answer using the information provided in the numbered sources below. Here's how:\"\n",
        "    \" - Read the existing answer and the sources carefully.\"\n",
        "    \" - Identify any information in the sources that can improve the answer by adding details, making it more accurate, or providing better support.\"\n",
        "    \" - If the sources provide new information, incorporate it into the answer.\"\n",
        "    \" - If the sources contradict the existing answer, correct the answer.\"\n",
        "    \" - If the sources aren't helpful, keep the original answer.\"\n",
        "    \"Cite the source number(s) for any information you use in your answer (e.g., [1]).\"\n",
        "    \"We have provided an existing answer: {existing_answer}\"\n",
        "    \"Below are several numbered sources of information. \"\n",
        "    \"Use them to refine the existing answer. \"\n",
        "    \"If the provided sources are not helpful, you will repeat the existing answer.\"\n",
        "    \"------\"\n",
        "    \"{context_msg}\"\n",
        "    \"------\"\n",
        "    \"Query: {query_str}\"\n",
        "    \"Answer: \"\n",
        ")\n",
        "\n",
        "DEFAULT_CITATION_CHUNK_SIZE = 512\n",
        "DEFAULT_CITATION_CHUNK_OVERLAP = 20"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "09009f1cd892"
      },
      "source": [
        "### Workflow Class\n",
        "\n",
        "The RAGWorkflow() class contains all the steps of the workflow. We define the steps by decorating the method with @step.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "486b4ba78947"
      },
      "outputs": [],
      "source": [
        "class RAGWorkflow(Workflow):\n",
        "    @step\n",
        "    async def ingest_data(\n",
        "        self, ctx: Context, ev: StartEvent\n",
        "    ) -> FirestoreIndexData | None:\n",
        "        \"\"\"Entry point to ingest a document, triggered by a StartEvent with 'dirname'.\"\"\"\n",
        "        dirname = ev.get(\"dirname\")\n",
        "        if not dirname:\n",
        "            return None\n",
        "\n",
        "        documents = SimpleDirectoryReader(dirname).load_data()\n",
        "        await ctx.set(\"documents\", documents)\n",
        "        return FirestoreIndexData(\n",
        "            status=\"First step complete. Data loaded into Documents.\"\n",
        "        )\n",
        "\n",
        "    @step\n",
        "    async def load_database(self, ctx: Context, ev: FirestoreIndexData) -> StopEvent:\n",
        "        print(ev.status)\n",
        "\n",
        "        # create (or load) docstore and add nodes\n",
        "        docstore = FirestoreDocumentStore.from_database(\n",
        "            project=PROJECT_ID,\n",
        "            database=FIRESTORE_DATABASE_ID,\n",
        "        )\n",
        "\n",
        "        docstore.add_documents(await ctx.get(\"documents\"))\n",
        "\n",
        "        # create storage context\n",
        "        storage_context = StorageContext.from_defaults(docstore=docstore)\n",
        "\n",
        "        # setup index\n",
        "        index = VectorStoreIndex.from_documents(\n",
        "            documents=await ctx.get(\"documents\"), storage_context=storage_context\n",
        "        )\n",
        "\n",
        "        print(\"Index created\")\n",
        "        return StopEvent(index)\n",
        "\n",
        "    def combine_queries(\n",
        "        self,\n",
        "        query_bundle: QueryBundle,\n",
        "        prev_reasoning: str,\n",
        "        llm: LLM,\n",
        "    ) -> QueryBundle:\n",
        "        \"\"\"Combine queries using StepDecomposeQueryTransform.\"\"\"\n",
        "        transform_metadata = {\"prev_reasoning\": prev_reasoning}\n",
        "        return StepDecomposeQueryTransform(llm=llm)(\n",
        "            query_bundle, metadata=transform_metadata\n",
        "        )\n",
        "\n",
        "    def default_stop_fn(self, stop_dict: dict) -> bool:\n",
        "        \"\"\"Stop function for multi-step query combiner.\"\"\"\n",
        "        query_bundle = cast(QueryBundle, stop_dict.get(\"query_bundle\"))\n",
        "        if query_bundle is None:\n",
        "            raise ValueError(\"Response must be provided to stop function.\")\n",
        "\n",
        "        return \"none\" in query_bundle.query_str.lower()\n",
        "\n",
        "    @step(pass_context=True)\n",
        "    async def query_multistep(\n",
        "        self, ctx: Context, ev: StartEvent\n",
        "    ) -> QueryMultiStepEvent | None:\n",
        "        \"\"\"Entry point for RAG, triggered by a StartEvent with `query`. Execute multi-step query process.\"\"\"\n",
        "\n",
        "        query = ev.get(\"query\")\n",
        "        index = ev.get(\"index\")\n",
        "\n",
        "        prev_reasoning = \"\"\n",
        "        cur_response = None\n",
        "        should_stop = False\n",
        "        cur_steps = 0\n",
        "\n",
        "        # use response\n",
        "        final_response_metadata: dict[str, Any] = {\"sub_qa\": []}\n",
        "\n",
        "        text_chunks = []\n",
        "        source_nodes = []\n",
        "\n",
        "        stop_fn = self.default_stop_fn\n",
        "\n",
        "        if not query:\n",
        "            return None\n",
        "\n",
        "        print(f\"Query the database with: {query}\")\n",
        "\n",
        "        # store the query in the global context\n",
        "        await ctx.set(\"query\", query)\n",
        "\n",
        "        # get the index from the global context\n",
        "        if index is None:\n",
        "            print(\"Index is empty, load some documents before querying!\")\n",
        "            return None\n",
        "\n",
        "        num_steps = ev.get(\"num_steps\")\n",
        "        query_engine = index.as_query_engine()\n",
        "\n",
        "        while not should_stop:\n",
        "            if num_steps is not None and cur_steps >= num_steps:\n",
        "                should_stop = True\n",
        "                break\n",
        "            elif should_stop:\n",
        "                break\n",
        "\n",
        "            updated_query_bundle = self.combine_queries(\n",
        "                QueryBundle(query_str=query),\n",
        "                prev_reasoning,\n",
        "                llm=Settings.llm,\n",
        "            )\n",
        "\n",
        "            print(\n",
        "                f\"Created query for the step - {cur_steps} is: {updated_query_bundle}\"\n",
        "            )\n",
        "\n",
        "            stop_dict = {\"query_bundle\": updated_query_bundle}\n",
        "            if stop_fn(stop_dict):\n",
        "                should_stop = True\n",
        "                break\n",
        "\n",
        "            cur_response = query_engine.query(updated_query_bundle)\n",
        "\n",
        "            # append to response builder\n",
        "            cur_qa_text = (\n",
        "                f\"\\nQuestion: {updated_query_bundle.query_str}\\n\"\n",
        "                f\"Answer: {cur_response!s}\"\n",
        "            )\n",
        "            text_chunks.append(cur_qa_text)\n",
        "            print(\"Source nodes used:\\n\")\n",
        "            for source_node in cur_response.source_nodes:\n",
        "                print(source_node)\n",
        "                source_nodes.append(source_node)\n",
        "\n",
        "            # update metadata\n",
        "            final_response_metadata[\"sub_qa\"].append(\n",
        "                (updated_query_bundle.query_str, cur_response)\n",
        "            )\n",
        "\n",
        "            prev_reasoning += (\n",
        "                f\"- {updated_query_bundle.query_str}\\n\" f\"- {cur_response!s}\\n\"\n",
        "            )\n",
        "            cur_steps += 1\n",
        "\n",
        "        nodes = [\n",
        "            NodeWithScore(node=TextNode(text=text_chunk)) for text_chunk in text_chunks\n",
        "        ]\n",
        "        return QueryMultiStepEvent(\n",
        "            nodes=nodes,\n",
        "            source_nodes=source_nodes,\n",
        "            final_response_metadata=final_response_metadata,\n",
        "        )\n",
        "\n",
        "    @step\n",
        "    async def rerank(self, ctx: Context, ev: QueryMultiStepEvent) -> RerankEvent:\n",
        "        # Rerank the nodes\n",
        "        ranker = LLMRerank(choice_batch_size=5, top_n=10, llm=Settings.llm)\n",
        "        print(\"Entering reranking of nodes:\\n\")\n",
        "        print(\"Original query: \", await ctx.get(\"query\", default=None), flush=True)\n",
        "        # print(await ctx.get(\"query\", default=None), flush=True)\n",
        "        try:\n",
        "            new_nodes = ranker.postprocess_nodes(\n",
        "                ev.nodes, query_str=await ctx.get(\"query\", default=None)\n",
        "            )\n",
        "        except:\n",
        "            # re ranker is not guaranteed to create parsable output\n",
        "            new_nodes = ev.nodes\n",
        "\n",
        "        print(f\"Reranked nodes to {len(new_nodes)}\")\n",
        "        return RerankEvent(\n",
        "            nodes=new_nodes,\n",
        "            source_nodes=ev.source_nodes,\n",
        "            final_response_metadata=ev.final_response_metadata,\n",
        "        )\n",
        "\n",
        "    @step\n",
        "    async def create_citation_nodes(self, ev: RerankEvent) -> CreateCitationsEvent:\n",
        "        \"\"\"\n",
        "        Modify retrieved nodes to create granular sources for citations.\n",
        "\n",
        "        Takes a list of NodeWithScore objects and splits their content\n",
        "        into smaller chunks, creating new NodeWithScore objects for each chunk.\n",
        "        Each new node is labeled as a numbered source, allowing for more precise\n",
        "        citation in query results.\n",
        "\n",
        "        Args:\n",
        "            nodes (List[NodeWithScore]): A list of NodeWithScore objects to be processed.\n",
        "\n",
        "        Returns:\n",
        "            List[NodeWithScore]: A new list of NodeWithScore objects, where each object\n",
        "            represents a smaller chunk of the original nodes, labeled as a source.\n",
        "        \"\"\"\n",
        "        nodes = ev.nodes\n",
        "\n",
        "        new_nodes: list[NodeWithScore] = []\n",
        "\n",
        "        text_splitter = SentenceSplitter(\n",
        "            chunk_size=DEFAULT_CITATION_CHUNK_SIZE,\n",
        "            chunk_overlap=DEFAULT_CITATION_CHUNK_OVERLAP,\n",
        "        )\n",
        "\n",
        "        for node in nodes:\n",
        "            print(node)\n",
        "\n",
        "            text_chunks = text_splitter.split_text(\n",
        "                node.node.get_content(metadata_mode=MetadataMode.NONE)\n",
        "            )\n",
        "\n",
        "            for text_chunk in text_chunks:\n",
        "                text = f\"Source {len(new_nodes)+1}:\\n{text_chunk}\\n\"\n",
        "\n",
        "                new_node = NodeWithScore(\n",
        "                    node=TextNode.model_validate(node.node), score=node.score\n",
        "                )\n",
        "\n",
        "                new_node.node.text = text\n",
        "                new_nodes.append(new_node)\n",
        "        return CreateCitationsEvent(\n",
        "            nodes=new_nodes,\n",
        "            source_nodes=ev.source_nodes,\n",
        "            final_response_metadata=ev.final_response_metadata,\n",
        "        )\n",
        "\n",
        "    @step\n",
        "    async def synthesize(self, ctx: Context, ev: CreateCitationsEvent) -> StopEvent:\n",
        "        \"\"\"Return a streaming response using reranked nodes.\"\"\"\n",
        "\n",
        "        print(\"Synthesizing final result...\")\n",
        "\n",
        "        response_synthesizer = get_response_synthesizer(\n",
        "            llm=Vertex(model=\"gemini-1.5-pro\", temperature=0.0, max_tokens=5000),\n",
        "            text_qa_template=CITATION_QA_TEMPLATE,\n",
        "            refine_template=CITATION_REFINE_TEMPLATE,\n",
        "            response_mode=ResponseMode.COMPACT,\n",
        "            use_async=True,\n",
        "        )\n",
        "        query = await ctx.get(\"query\", default=None)\n",
        "        response = await response_synthesizer.asynthesize(\n",
        "            query, nodes=ev.nodes, additional_source_nodes=ev.source_nodes\n",
        "        )\n",
        "        return StopEvent(result=response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "506c0759eab2"
      },
      "outputs": [],
      "source": [
        "# optional - generate DAG for workflow created above\n",
        "draw_all_possible_flows(workflow=RAGWorkflow, filename=\"multi_step_workflow.html\")  # type: ignore"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4eebc39a2b70"
      },
      "source": [
        "### Run the workflow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "d7a98c8d7876"
      },
      "outputs": [],
      "source": [
        "w = RAGWorkflow(timeout=200)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "0a14595e4e6a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "First step complete. Data loaded into Documents.\n",
            "Index created\n"
          ]
        }
      ],
      "source": [
        "# Ingest the documents\n",
        "index = await w.run(dirname=\"./data\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "439d69a3e7f7"
      },
      "source": [
        "#### Example 1\n",
        "Query: \"What is the significance of the green light?\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "19ebb8696f71"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Query the database with: What is the significance of the green light?\n",
            "Created query for the step - 0 is: What is the significance of the green light?\n",
            "Source nodes used:\n",
            "\n",
            "Node ID: 0eab96dd-33ef-4d5c-a97e-8ca897af48d6\n",
            "Text: Its vanished trees, the trees that had made way for Gatsby’s\n",
            "house, had once pandered in whispers to the last and greatest of all\n",
            "human dreams; for a transitory enchanted moment man must have held his\n",
            "breath in the presence of this continent, compelled into an aesthetic\n",
            "contemplation he neither understood nor desired, face to face for the\n",
            "l...\n",
            "Score:  0.540\n",
            "\n",
            "Node ID: 4b08ce92-cbf0-4469-88a5-8cb3514da22f\n",
            "Text: “I’ve got a man in England who buys me clothes. He sends over a\n",
            "selection of things at the beginning of each season, spring and fall.”\n",
            "He took out a pile of shirts and began throwing them, one by one,\n",
            "before us, shirts of sheer linen and thick silk and fine flannel,\n",
            "which lost their folds as they fell and covered the table in  many-\n",
            "coloure...\n",
            "Score:  0.525\n",
            "\n",
            "Created query for the step - 1 is: ## New Question:\n",
            "\n",
            "**What is the significance of the green light in the context of Gatsby's pursuit of Daisy?** \n",
            "\n",
            "Source nodes used:\n",
            "\n",
            "Node ID: f323395e-7546-454a-9f8b-563e73fbb292\n",
            "Text: “Old sport, the dance is unimportant.”    He wanted nothing less\n",
            "of Daisy than that she should go to Tom and  say: “I never loved you.”\n",
            "After she had obliterated four years with  that sentence they could\n",
            "decide upon the more practical measures to be  taken. One of them was\n",
            "that, after she was free, they were to go back  to Louisville and be\n",
            "marr...\n",
            "Score:  0.662\n",
            "\n",
            "Node ID: a2ec7e02-2983-4da9-b08a-afa1b6cc4216\n",
            "Text: “Why didn’t he ask you to arrange a meeting?”    “He wants her\n",
            "to see his house,” she explained. “And your house is  right next\n",
            "door.”    “Oh!”    “I think he half expected her to wander into one of\n",
            "his parties, some  night,” went on Jordan, “but she never did. Then he\n",
            "began asking  people casually if they knew her, and I was the first\n",
            "one he fo...\n",
            "Score:  0.648\n",
            "\n",
            "Entering reranking of nodes:\n",
            "\n",
            "Original query:  What is the significance of the green light?\n",
            "Reranked nodes to 2\n",
            "Node ID: c2860521-c9c1-4cab-b7a9-ea1c784506be\n",
            "Text: Question: What is the significance of the green light? Answer:\n",
            "The green light is a symbol of Gatsby's dream of Daisy. It is the\n",
            "light at the end of her dock, which he can see from his house across\n",
            "the bay. The green light represents Gatsby's hope for a future with\n",
            "Daisy, and his belief that he can recapture the past. However, the\n",
            "green light is...\n",
            "Score: None\n",
            "\n",
            "Node ID: 7fe78bba-c870-486e-8f29-0168b09a792e\n",
            "Text: Question: ## New Question:  **What is the significance of the\n",
            "green light in the context of Gatsby's pursuit of Daisy?**   Answer:\n",
            "## The Green Light: A Symbol of Gatsby's Dreams and Desires  The green\n",
            "light at the end of Daisy's dock plays a pivotal role in symbolizing\n",
            "Gatsby's aspirations and the unattainable nature of his dreams. It\n",
            "represent...\n",
            "Score: None\n",
            "\n",
            "Synthesizing final result...\n"
          ]
        },
        {
          "data": {
            "text/markdown": [
              "## The Significance of the Green Light in The Great Gatsby\n",
              "\n",
              "The green light at the end of Daisy's dock holds immense symbolic weight in F. Scott Fitzgerald's *The Great Gatsby*. It represents a multitude of Gatsby's aspirations and desires, while simultaneously highlighting the unattainable nature of his dreams.\n",
              "\n",
              "**Unrequited Love:** The green light's physical proximity to Gatsby, yet separation by the bay, mirrors the emotional distance between him and Daisy. He yearns for her, but she remains out of reach, symbolizing his unrequited love.\n",
              "\n",
              "**The Past:** The green light evokes memories of Gatsby's past with Daisy, a time when their love seemed possible. He desperately wants to recapture that lost time and recreate their romance, clinging to the hope of a second chance.\n",
              "\n",
              "**Hope and Illusion:** The green light embodies Gatsby's unwavering hope for a future with Daisy. He believes that if he can achieve enough wealth and success, he can win her back. However, this hope is ultimately an illusion, as Daisy has moved on and their circumstances have changed.\n",
              "\n",
              "**The American Dream:** The green light can be interpreted as a symbol of the American Dream, representing Gatsby's relentless pursuit of wealth and social status. He believes that achieving these goals will bring him happiness and allow him to win Daisy's love. However, the novel ultimately suggests that the American Dream is often unattainable and can lead to disillusionment.\n",
              "\n",
              "**Additional Points:**\n",
              "\n",
              "* The green light's color reinforces its symbolic meaning. Green often represents hope, growth, and new beginnings, but in this context, it takes on a more melancholic and unattainable quality.\n",
              "* The light's flickering nature reflects the instability of Gatsby's dreams and the uncertainty of his future.\n",
              "* Gatsby's constant focus on the green light highlights his single-minded obsession with Daisy and his inability to move on from the past.\n",
              "\n",
              "**Overall, the green light serves as a powerful symbol that encapsulates Gatsby's longing, his yearning for a lost love, and the ultimately unattainable nature of his dreams.**\n",
              "\n",
              "**Sources:**\n",
              "\n",
              "* [1] The Great Gatsby by F. Scott Fitzgerald\n",
              "* [2] SparkNotes: The Great Gatsby - Symbols, Imagery, Allegory"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# Run a query\n",
        "NUM_STEPS = 2  # @param {type:\"int\"} represents how many sub-questions generated based on the query\n",
        "result = await w.run(\n",
        "    query=\"What is the significance of the green light?\",\n",
        "    index=index,\n",
        "    num_steps=NUM_STEPS,\n",
        ")\n",
        "\n",
        "display(Markdown(f\"{result}\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "da168aa48f73"
      },
      "source": [
        "Check the ranked LLM generated sub-question answers used:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "5dd8dab92106"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Node ID: c2860521-c9c1-4cab-b7a9-ea1c784506be\n",
            "Text: Source 1: Question: What is the significance of the green light?\n",
            "Answer: The green light is a symbol of Gatsby's dream of Daisy. It is\n",
            "the light at the end of her dock, which he can see from his house\n",
            "across the bay. The green light represents Gatsby's hope for a future\n",
            "with Daisy, and his belief that he can recapture the past. However,\n",
            "the gree...\n",
            "Score: None\n",
            "\n",
            "Node ID: 7fe78bba-c870-486e-8f29-0168b09a792e\n",
            "Text: Source 2: Question: ## New Question:  **What is the significance\n",
            "of the green light in the context of Gatsby's pursuit of Daisy?**\n",
            "Answer: ## The Green Light: A Symbol of Gatsby's Dreams and Desires\n",
            "The green light at the end of Daisy's dock plays a pivotal role in\n",
            "symbolizing Gatsby's aspirations and the unattainable nature of his\n",
            "dreams. It...\n",
            "Score: None\n",
            "\n"
          ]
        }
      ],
      "source": [
        "for idx in range(0, NUM_STEPS):\n",
        "    print(result.source_nodes[idx])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "79eb31af5944"
      },
      "source": [
        "Check the citations from the original source used:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "6ddf240120c0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Node ID: 0eab96dd-33ef-4d5c-a97e-8ca897af48d6\n",
            "Text: Its vanished trees, the trees that had made way for Gatsby’s\n",
            "house, had once pandered in whispers to the last and greatest of all\n",
            "human dreams; for a transitory enchanted moment man must have held his\n",
            "breath in the presence of this continent, compelled into an aesthetic\n",
            "contemplation he neither understood nor desired, face to face for the\n",
            "l...\n",
            "Score:  0.540\n",
            "\n",
            "Node ID: 4b08ce92-cbf0-4469-88a5-8cb3514da22f\n",
            "Text: “I’ve got a man in England who buys me clothes. He sends over a\n",
            "selection of things at the beginning of each season, spring and fall.”\n",
            "He took out a pile of shirts and began throwing them, one by one,\n",
            "before us, shirts of sheer linen and thick silk and fine flannel,\n",
            "which lost their folds as they fell and covered the table in  many-\n",
            "coloure...\n",
            "Score:  0.525\n",
            "\n",
            "Node ID: f323395e-7546-454a-9f8b-563e73fbb292\n",
            "Text: “Old sport, the dance is unimportant.”    He wanted nothing less\n",
            "of Daisy than that she should go to Tom and  say: “I never loved you.”\n",
            "After she had obliterated four years with  that sentence they could\n",
            "decide upon the more practical measures to be  taken. One of them was\n",
            "that, after she was free, they were to go back  to Louisville and be\n",
            "marr...\n",
            "Score:  0.662\n",
            "\n",
            "Node ID: a2ec7e02-2983-4da9-b08a-afa1b6cc4216\n",
            "Text: “Why didn’t he ask you to arrange a meeting?”    “He wants her\n",
            "to see his house,” she explained. “And your house is  right next\n",
            "door.”    “Oh!”    “I think he half expected her to wander into one of\n",
            "his parties, some  night,” went on Jordan, “but she never did. Then he\n",
            "began asking  people casually if they knew her, and I was the first\n",
            "one he fo...\n",
            "Score:  0.648\n",
            "\n"
          ]
        }
      ],
      "source": [
        "for idx in range(NUM_STEPS, len(result.source_nodes)):\n",
        "    print(result.source_nodes[idx])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2dca7a130be9"
      },
      "source": [
        "## Cleaning up\n",
        "\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_workflows.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}