{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "57ec3ba4dca6"
      },
      "source": [
        "# Small-to-big Retrieval-Augmented Generation  \n",
        "\n",
        "<table align=\"left\">\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%2Fsmall_to_big_rag%2Fsmall_to_big_rag.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\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/small_to_big_rag/small_to_big_rag.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Run in Colab\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/small_to_big_rag/small_to_big_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/small_to_big_rag/small_to_big_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/small_to_big_rag/small_to_big_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/small_to_big_rag/small_to_big_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/small_to_big_rag/small_to_big_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/small_to_big_rag/small_to_big_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/small_to_big_rag/small_to_big_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": "b130a2a11d10"
      },
      "source": [
        "| Authors |\n",
        "| --- |\n",
        "| [Keith Ballinger](https://github.com/keithballinger) |\n",
        "| [Megan O'Keefe](https://github.com/askmeegs) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "603dbcc4096e"
      },
      "source": [
        "## Overview\n",
        "\n",
        "**YouTube Video: Using RAG expansion to improve model speed and accuracy**\n",
        "\n",
        "<a href=\"https://www.youtube.com/watch?v=I_l91kdNE4w&list=PLIivdWyY5sqLvGdVLJZh2EMax97_T-OIB\" target=\"_blank\">\n",
        "  <img src=\"https://img.youtube.com/vi/I_l91kdNE4w/maxresdefault.jpg\" alt=\"Using RAG expansion to improve model speed and accuracy\" width=\"500\">\n",
        "</a>\n",
        "\n",
        "Small-to-big retrieval is a form of [modular recursive RAG](https://www.promptingguide.ai/research/rag#modular-rag), where you link smaller grounding data chunks to larger \"parent\" data chunks. When a small chunk is retrieved at runtime, the larger linked chunk can be retrieved if needed.  \n",
        "\n",
        "The Small-to-big strategy offers a few benefits over regular RAG: \n",
        "1. **Complex use cases**: Small-to-big RAG can be used to handle complex queries where the context is too large to fit into a single dense vector. Examples: legal documents, research papers. \n",
        "2. **Work around the limits of dense vectors**: Dense vectors can only \"squish\" the meaning of a text so much. If you try to embed a chunk that's too long (eg. an entire document), some of the meaning may be lost, resulting in less accurate retrieval results. By keeping the small vector chunks small, and retrieving large documents later (with or without embeddings), you're getting the semantic-search benefits of dense vectors, while still being able to retrieve the full context when needed.\n",
        "3. **Cost**: You can set up small-to-big RAG to only fetch the long documents when needed (eg. if the model is unable to respond with the small context). This can save on inference costs, because Gemini on Vertex AI is [priced per input character](https://cloud.google.com/vertex-ai/generative-ai/pricing). \n",
        "\n",
        "There are multiple ways to implement small-to-big RAG. The small chunks could represent short passages of a document, and the larger chunks could represent the entire surrounding context (eg. the whole document) - see LangChain's [ParentDocumentRetriever](https://python.langchain.com/docs/modules/data_connection/retrievers/parent_document_retriever) for more info. Another way is to have the small chunks represent summaries of the larger documents. That is the method we'll explore here. \n",
        "\n",
        "In this example, we'll walk through a Small-to-big RAG example using a GitHub codebase called [Online Boutique](https://github.com/GoogleCloudPlatform/microservices-demo). Online Boutique is a microservices, multi-language sample application. We'll implement a question answering functionality to help a new contributor learn about and navigate this codebase."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9b8267f065f8"
      },
      "source": [
        "![](architecture.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e7c7de7a2c56"
      },
      "source": [
        "To complete this notebook, **you will need**: \n",
        "- A [Google Cloud account](https://console.cloud.google.com/)\n",
        "- One [Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects) with [billing](https://cloud.google.com/billing/docs/how-to/modify-project) enabled \n",
        "- Enable the [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)\n",
        "- [gcloud SDK](https://cloud.google.com/sdk/docs/install) installed in your environment.\n",
        "- Your user has the **Vertex AI User** IAM role.\n",
        "\n",
        "This notebook uses the following products and tools:\n",
        "- [Vertex AI - Gemini API](https://cloud.google.com/vertex-ai/generative-ai/docs/start/quickstarts/quickstart-multimodal#gemini-beginner-samples-python_vertex_ai_sdk) \n",
        "- [Vertex AI - Text Embeddings API](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings) \n",
        "- [Chroma](https://docs.trychroma.com/getting-started) (in-memory vector database)  \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f55c103999de"
      },
      "source": [
        "### Setup "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f30a73cd08af"
      },
      "source": [
        "First, set variables. You are required to set your project ID. You can keep the other variables as-is."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "709610997696"
      },
      "outputs": [],
      "source": [
        "EMBEDDING_MODEL = \"text-embedding-005\"\n",
        "GENERATIVE_MODEL = \"gemini-2.0-flash\"\n",
        "PROJECT_ID = \"YOUR-PROJECT-ID\"\n",
        "REGION = \"us-central1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "be8bd58227e6"
      },
      "source": [
        "Install the necessary packages, and import them. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d24ab152f767"
      },
      "outputs": [],
      "source": [
        "%pip install google-cloud-aiplatform pandas chromadb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "7947eb28a61c"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import chromadb\n",
        "import pandas as pd\n",
        "import vertexai\n",
        "from vertexai.generative_models import ChatSession, GenerativeModel\n",
        "from vertexai.language_models import TextEmbeddingModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cdc1ba9b8efe"
      },
      "source": [
        "Lastly, download the source code dataset from Cloud Storage. This is a modified version of the upstream Online Boutique repo, with certain files pruned for sample purposes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "701c9ae117c1"
      },
      "outputs": [],
      "source": [
        "! gsutil -m cp -r gs://github-repo/generative-ai/gemini/use-cases/rag/small-to-big-rag/onlineboutique-codefiles ."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c341984c017d"
      },
      "source": [
        "### Create helper functions\n",
        "\n",
        "We'll create one function that calls Vertex AI text embeddings, and another that inferences Gemini on Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "6289a6934a09"
      },
      "outputs": [],
      "source": [
        "model = TextEmbeddingModel.from_pretrained(EMBEDDING_MODEL)\n",
        "\n",
        "\n",
        "def get_text_embedding(doc) -> list:\n",
        "    embeddings = model.get_embeddings([doc])\n",
        "    if len(embeddings) > 1:\n",
        "        raise ValueError(\"More than one embedding returned.\")\n",
        "    if len(embeddings) == 0:\n",
        "        raise ValueError(\"No embedding returned.\")\n",
        "    return embeddings[0].values"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "fe5b30c3e3c0"
      },
      "outputs": [],
      "source": [
        "vertexai.init(project=PROJECT_ID, location=REGION)\n",
        "model = GenerativeModel(GENERATIVE_MODEL)\n",
        "chat = model.start_chat()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "2809ed0a38d3"
      },
      "outputs": [],
      "source": [
        "def gemini_inference(chat: ChatSession, prompt: str) -> str:\n",
        "    text_response = chat.generate_content(prompt)\n",
        "    return text_response.text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "91bf6770a06b"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Bonjour le monde!'"
            ]
          },
          "execution_count": 26,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "gemini_inference(chat, \"hello world!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "361663cfb9ff"
      },
      "source": [
        "### Get summaries of code files \n",
        "\n",
        "First, we'll use Gemini on Vertex AI to get short summaries of each code file.  We'll do this by recursively traversing the files in `onlineboutique-codefiles`. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6cde1cce5ef6"
      },
      "outputs": [],
      "source": [
        "# for every file in onlineboutique-codefiles/, read it in, and get the full tree filename, and a code summary\n",
        "summaries = {}\n",
        "for root, dirs, files in os.walk(\"onlineboutique-codefiles/\"):\n",
        "    for file in files:\n",
        "        temp = {}\n",
        "        full_file_path = os.path.join(root, file)\n",
        "        with open(full_file_path) as f:\n",
        "            print(\"Processing file: \", full_file_path)\n",
        "            try:\n",
        "                content = f.read()\n",
        "                temp[\"content\"] = content\n",
        "                prompt = \"\"\" \n",
        "                You are a helpful code summarizer. Here is a source code file. Please identify the programming language and summarize it in three sentences or less. Give as much detail as possible, including function names and libraries used. Code: \n",
        "                {}\n",
        "                \"\"\".format(\n",
        "                    content\n",
        "                )\n",
        "                summary = gemini_inference(chat, prompt)\n",
        "                temp[\"summary\"] = summary\n",
        "                summaries[full_file_path] = temp\n",
        "            except Exception as e:\n",
        "                print(f\"⚠️ Error processing file: {full_file_path} - {e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "075419a1669e"
      },
      "source": [
        "Next, we'll create a Pandas DataFrame with the file paths, code content, and summaries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "6bb17c215f80"
      },
      "outputs": [],
      "source": [
        "df = pd.DataFrame.from_dict(summaries, orient=\"index\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "id": "06d0aaec00df"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>content</th>\n",
              "      <th>summary</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>onlineboutique-codefiles/LICENSE</th>\n",
              "      <td>\\n                                 Apache Lice...</td>\n",
              "      <td>This is a text file describing the Apache Lice...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>onlineboutique-codefiles/cloudbuild.yaml</th>\n",
              "      <td># Copyright 2020 Google LLC\\n#\\n# Licensed und...</td>\n",
              "      <td>This configuration file uses Google Cloud Buil...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>onlineboutique-codefiles/README.md</th>\n",
              "      <td>&lt;p align=\"center\"&gt;\\n&lt;img src=\"/src/frontend/st...</td>\n",
              "      <td>Online Boutique is a web-based microservices d...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>onlineboutique-codefiles/skaffold.yaml</th>\n",
              "      <td># Copyright 2021 Google LLC\\n#\\n# Licensed und...</td>\n",
              "      <td>This Skaffold configuration file defines build...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>onlineboutique-codefiles/terraform/output.tf</th>\n",
              "      <td># Copyright 2022 Google LLC\\n#\\n# Licensed und...</td>\n",
              "      <td>This Terraform configuration outputs data from...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                                                                                        content  \\\n",
              "onlineboutique-codefiles/LICENSE              \\n                                 Apache Lice...   \n",
              "onlineboutique-codefiles/cloudbuild.yaml      # Copyright 2020 Google LLC\\n#\\n# Licensed und...   \n",
              "onlineboutique-codefiles/README.md            <p align=\"center\">\\n<img src=\"/src/frontend/st...   \n",
              "onlineboutique-codefiles/skaffold.yaml        # Copyright 2021 Google LLC\\n#\\n# Licensed und...   \n",
              "onlineboutique-codefiles/terraform/output.tf  # Copyright 2022 Google LLC\\n#\\n# Licensed und...   \n",
              "\n",
              "                                                                                        summary  \n",
              "onlineboutique-codefiles/LICENSE              This is a text file describing the Apache Lice...  \n",
              "onlineboutique-codefiles/cloudbuild.yaml      This configuration file uses Google Cloud Buil...  \n",
              "onlineboutique-codefiles/README.md            Online Boutique is a web-based microservices d...  \n",
              "onlineboutique-codefiles/skaffold.yaml        This Skaffold configuration file defines build...  \n",
              "onlineboutique-codefiles/terraform/output.tf  This Terraform configuration outputs data from...  "
            ]
          },
          "execution_count": 29,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "df.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "id": "cb7b939dd1b3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Number of rows:  109\n"
          ]
        }
      ],
      "source": [
        "# number of file summaries\n",
        "print(\"Number of rows: \", df.shape[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "id": "7c795eb2c3ce"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>filename</th>\n",
              "      <th>content</th>\n",
              "      <th>summary</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>onlineboutique-codefiles/LICENSE</td>\n",
              "      <td>\\n                                 Apache Lice...</td>\n",
              "      <td>This is a text file describing the Apache Lice...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>onlineboutique-codefiles/cloudbuild.yaml</td>\n",
              "      <td># Copyright 2020 Google LLC\\n#\\n# Licensed und...</td>\n",
              "      <td>This configuration file uses Google Cloud Buil...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>onlineboutique-codefiles/README.md</td>\n",
              "      <td>&lt;p align=\"center\"&gt;\\n&lt;img src=\"/src/frontend/st...</td>\n",
              "      <td>Online Boutique is a web-based microservices d...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>onlineboutique-codefiles/skaffold.yaml</td>\n",
              "      <td># Copyright 2021 Google LLC\\n#\\n# Licensed und...</td>\n",
              "      <td>This Skaffold configuration file defines build...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>onlineboutique-codefiles/terraform/output.tf</td>\n",
              "      <td># Copyright 2022 Google LLC\\n#\\n# Licensed und...</td>\n",
              "      <td>This Terraform configuration outputs data from...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                                       filename  \\\n",
              "0              onlineboutique-codefiles/LICENSE   \n",
              "1      onlineboutique-codefiles/cloudbuild.yaml   \n",
              "2            onlineboutique-codefiles/README.md   \n",
              "3        onlineboutique-codefiles/skaffold.yaml   \n",
              "4  onlineboutique-codefiles/terraform/output.tf   \n",
              "\n",
              "                                             content  \\\n",
              "0  \\n                                 Apache Lice...   \n",
              "1  # Copyright 2020 Google LLC\\n#\\n# Licensed und...   \n",
              "2  <p align=\"center\">\\n<img src=\"/src/frontend/st...   \n",
              "3  # Copyright 2021 Google LLC\\n#\\n# Licensed und...   \n",
              "4  # Copyright 2022 Google LLC\\n#\\n# Licensed und...   \n",
              "\n",
              "                                             summary  \n",
              "0  This is a text file describing the Apache Lice...  \n",
              "1  This configuration file uses Google Cloud Buil...  \n",
              "2  Online Boutique is a web-based microservices d...  \n",
              "3  This Skaffold configuration file defines build...  \n",
              "4  This Terraform configuration outputs data from...  "
            ]
          },
          "execution_count": 32,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# the first column should be named \"filename\"\n",
        "df = df.reset_index()\n",
        "df = df.rename(columns={\"index\": \"filename\"})\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "id": "d9fd2ea75163"
      },
      "outputs": [],
      "source": [
        "# write to csv\n",
        "df.to_csv(\"code_summaries.csv\", index=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f5411d82c592"
      },
      "source": [
        "### Convert summaries to embeddings\n",
        "\n",
        "Next, we'll convert the text summaries of each code file to vector embeddings. We'll store those embeddings in an in-memory Chroma database. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "a15f516cdaa2"
      },
      "outputs": [],
      "source": [
        "chroma_client = chromadb.Client()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "fcacd266d4e6"
      },
      "outputs": [],
      "source": [
        "collection = chroma_client.create_collection(name=\"code_summaries\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aa51c5c900d6"
      },
      "outputs": [],
      "source": [
        "# iterate over dataframe. convert summary into embeddings. insert summary into collection.\n",
        "for index, row in df.iterrows():\n",
        "    fn = row[\"filename\"]\n",
        "    print(\"Getting embedding for: \", fn)\n",
        "    summary = row[\"summary\"]\n",
        "    print(summary)\n",
        "    e = get_text_embedding(summary)\n",
        "    print(e)\n",
        "    # add vector embedding to in-memory Chroma database.\n",
        "    # the \"small\" summary embedding is linked to the \"big\" raw code file through the metadata key, \"filename.\"\n",
        "    collection.add(\n",
        "        embeddings=[e], documents=[summary], metadatas=[{\"filename\": fn}], ids=[fn]\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "323b0b72a63d"
      },
      "source": [
        "### Implement the Small-to-big RAG workflow "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b2dc395f3af3"
      },
      "outputs": [],
      "source": [
        "# Get a list of all files to pass to Gemini, if it needs to see a specific code file.\n",
        "all_files = []\n",
        "for root, dirs, files in os.walk(\"onlineboutique-codefiles/\"):\n",
        "    for file in files:\n",
        "        all_files.append(os.path.join(root, file))\n",
        "print(all_files)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "51cf5003f530"
      },
      "source": [
        "The function below shows how we'll first try to inference Gemini with small chunks (code file summaries). If Gemini can answer with that context, we return its response and we're done. If Gemini needs more context, we'll ask it what file it would like to see. Then, we'll directly retrieve the code file from the DataFrame, and pass it into Gemini again as the \"large\" context."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 63,
      "metadata": {
        "id": "d2f0f9237755"
      },
      "outputs": [],
      "source": [
        "def small_to_big(user_prompt):\n",
        "    # SMALL: first, run RAG with the summary embeddings to try to get a response\n",
        "    query_emb = get_text_embedding(user_prompt)\n",
        "    result = collection.query(query_embeddings=[query_emb], n_results=3)\n",
        "    # process nearest-neighbors\n",
        "    processed_result = {}\n",
        "    d = result[\"documents\"][0]\n",
        "    for i in range(0, len(d)):\n",
        "        summary = d[i]\n",
        "        filename = result[\"metadatas\"][0][i][\"filename\"]\n",
        "        processed_result[filename] = summary\n",
        "    prompt_with_small = \"\"\"\n",
        "    You are a codebase helper. You will be given a user's question about the codebase, along with \n",
        "    summaries of relevant code files. Attempt to answer the question and only respond if you're confident in the answer. \n",
        "    If you need any more information, respond with ONLY the phrase \"need more context\". \n",
        "\n",
        "    The user query is: {} \n",
        "\n",
        "    The summaries are: {}\n",
        "    \"\"\".format(\n",
        "        user_prompt, str(processed_result)\n",
        "    )\n",
        "    print(prompt_with_small)\n",
        "    small_result = gemini_inference(chat, prompt_with_small)\n",
        "    # we're done if Gemini is confident with just the summaries as context...\n",
        "    if \"need more context\" not in small_result.lower():\n",
        "        return (\n",
        "            \"🐝 Completed at small, Gemini had enough context to respond. RESPONSE: \\n\"\n",
        "            + small_result\n",
        "        )\n",
        "    print(\n",
        "        \"🤔 Gemini asked for more context. Let's ask what codefile Gemini wants to see.\"\n",
        "    )\n",
        "    # otherwise, move on to BIG:\n",
        "    # IF we need the full context, get the filename that most closely matches the user's question\n",
        "    prompt_to_get_filename = \"\"\" \n",
        "    You are a codebase helper. The list of code files that you know about: \n",
        "    {}\n",
        "\n",
        "    The user asks the following question about the codebase: {}\n",
        "\n",
        "    Please respond with the filename that most closely matches the user's question. Respond with ONLY the filename. \n",
        "    \"\"\".format(\n",
        "        all_files, user_prompt\n",
        "    )\n",
        "    filename = gemini_inference(chat, prompt_to_get_filename)\n",
        "    print(\"📂 Gemini asked for this file: \" + filename)\n",
        "    # is the filename in the dataframe?\n",
        "    if filename not in df[\"filename\"].values:\n",
        "        # attempt to try again, appending \"onlineboutique-codefiles\"\n",
        "        filename = \"onlineboutique-codefiles/\" + filename\n",
        "        if filename not in df[\"filename\"].values:\n",
        "            return f\"⚠️ Error: filename {filename} not found in dataframe\"\n",
        "\n",
        "    # get the full code file\n",
        "    full_code = df[df[\"filename\"] == filename][\"content\"].values[0]\n",
        "    prompt_with_big = \"\"\" \n",
        "    You are a codebase helper. You will be given a user's question about the codebase, along with a complete source code file. Respond to the user's question with as much detail as possible.\n",
        "\n",
        "    The user query is: {}\n",
        "    \n",
        "    The full code file is: {}\n",
        "    \"\"\".format(\n",
        "        user_prompt, full_code\n",
        "    )\n",
        "\n",
        "    big_response = gemini_inference(chat, prompt_with_big)\n",
        "    return \"🦖 Completed at big. RESPONSE: \\n\" + big_response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "435ea72e8f3d"
      },
      "source": [
        "### Test it out \n",
        "\n",
        "You can test this function by calling `small_to_big(\"your codebase question\")`. We include a few examples below. The more detailed your question (eg. citing a specific function, line of code, or dependency), the more likely it is that Gemini will ask for more context, and the \"large\" step occurs. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {
        "id": "fed74a7235a3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "    You are a codebase helper. You will be given a user's question about the codebase, along with \n",
            "    summaries of relevant code files. Attempt to answer the question and only respond if you're confident in the answer. \n",
            "    If you need any more information, respond with ONLY the phrase \"need more context\". \n",
            "\n",
            "    The user query is: How does the ad service work? \n",
            "\n",
            "    The summaries are: {'onlineboutique-codefiles/src/adservice/src/main/java/hipstershop/AdService.java': 'This is a Java gRPC server for an ad service.\\nIt has a `getAds` function that takes a request containing context and returns a response with a list of ads.\\nIt uses a map of categories to ads to retrieve the ads and a random function to retrieve random ads if no category is specified.', 'onlineboutique-codefiles/src/adservice/src/main/java/hipstershop/AdServiceClient.java': 'This Java code defines a client that interacts with an Ads Service using gRPC. The `AdServiceClient` class establishes a connection to the service running at a specified address and port. Once connected, it sends `getAds` requests with a context key and receives `AdResponse` objects containing relevant advertisements. The client uses `ManagedChannel` for communication and `AdServiceGrpc.AdServiceBlockingStub` to invoke RPC methods. It also utilizes the `log4j2` library for logging.', 'onlineboutique-codefiles/src/adservice/README.md': 'This code is written in Java and uses the Gradle build tool to compile and package the application. The code includes a simple service called \"AdService\" that provides advertisements based on context keys. The service can be built locally using the `./gradlew installDist` command or by building a Docker image using the `docker build ./` command from the `src/adservice/` directory.'}\n",
            "    \n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "'🐝 Completed at small, Gemini had enough context to respond. RESPONSE: \\nThe Ad Service is a gRPC service that provides advertisements based on context keys. It has a `getAds` function that takes a request containing a context key and returns a response with a list of ads.\\nThe service uses a map of categories to ads to retrieve the ads, and a random function to retrieve random ads if no category is specified.\\nThe Ads Service is implemented in Java and uses the Gradle build tool to compile and package the application. It can be built locally using the `./gradlew installDist` command or by building a Docker image using the `docker build ./` command from the `src/adservice/` directory.'"
            ]
          },
          "execution_count": 55,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# an example of a query where only the small (summary) step is needed\n",
        "small_to_big(\"How does the ad service work?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 65,
      "metadata": {
        "id": "a15db3597709"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "    You are a codebase helper. You will be given a user's question about the codebase, along with \n",
            "    summaries of relevant code files. Attempt to answer the question and only respond if you're confident in the answer. \n",
            "    If you need any more information, respond with ONLY the phrase \"need more context\". \n",
            "\n",
            "    The user query is: Exactly how long is the kubectl wait condition in the Terraform deployment of online boutique? Return the right number of seconds \n",
            "\n",
            "    The summaries are: {'onlineboutique-codefiles/terraform/README.md': \"This Terraform script deploys the Online Boutique sample application, a microservices-based application, on a Google Kubernetes Engine (GKE) cluster. It provisions resources such as a GKE cluster, Memorystore (Redis) instance (optional), and a Kubernetes deployment for the application. Functions like `terraform init` and `terraform apply` are used to initialize and create the resources. Libraries like `kubectl` are used to retrieve the frontend's external IP address for accessing the application.\", 'onlineboutique-codefiles/terraform/main.tf': 'This Terraform code defines an infrastructure using Cloud APIs, a GKE cluster, and kubectl commands. It starts by defining local variables and modules to enable Google services. It then sets up a GKE cluster with autopilot enabled. Following this, it retrieves cluster credentials, and finally uses kubectl to deploy YAML configurations and waits until Pods are ready. The code relies on the `google` and `terraform-google-modules/gcloud` modules.', 'onlineboutique-codefiles/terraform/output.tf': 'This Terraform configuration outputs data from a Google Container Cluster resource. The output includes the `location` and `name` of the cluster. The `location` variable is assigned the value of the `location` attribute of the `resource.google_container_cluster.my_cluster` object, and the `name` variable is assigned the value of the `name` attribute of the same object.'}\n",
            "    \n",
            "🤔 Gemini asked for more context. Let's ask what codefile Gemini wants to see.\n",
            "📂 Gemini asked for this file: terraform/main.tf\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "'🦖 Completed at big. RESPONSE: \\nThe kubectl wait condition in the Terraform deployment of online boutique waits for two conditions to be met before finishing:\\n\\n1. The `apiservice/v1beta1.metrics.k8s.io` API service to be available.\\n2. All pods in the specified namespace to be ready.\\n\\nThe `--timeout` parameter specifies the maximum amount of time to wait for each condition to be met. In this case, the timeout is set to 180 seconds for the API service and 280 seconds for the pods.\\n\\nTherefore, the right number of seconds to use in the `kubectl wait condition` command is either 180 or 280, depending on which condition takes longer to be met.'"
            ]
          },
          "execution_count": 65,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "small_to_big(\n",
        "    \"Exactly how long is the kubectl wait condition in the Terraform deployment of online boutique? Return the right number of seconds\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b43a6cb2ca35"
      },
      "outputs": [],
      "source": [
        "# Solution terraform code in main.tf  - 280 seconds is correct\n",
        "\"\"\"\n",
        "\n",
        "# Wait condition for all Pods to be ready before finishing\n",
        "resource \"null_resource\" \"wait_conditions\" {\n",
        "  provisioner \"local-exec\" {\n",
        "    interpreter = [\"bash\", \"-exc\"]\n",
        "    command     = <<-EOT\n",
        "    kubectl wait --for=condition=AVAILABLE apiservice/v1beta1.metrics.k8s.io --timeout=180s\n",
        "    kubectl wait --for=condition=ready pods --all -n ${var.namespace} --timeout=280s\n",
        "    EOT\n",
        "  }\n",
        "\n",
        "  depends_on = [\n",
        "    resource.null_resource.apply_deployment\n",
        "  ]\n",
        "}\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 61,
      "metadata": {
        "id": "6f84a666ad01"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "    You are a codebase helper. You will be given a user's question about the codebase, along with \n",
            "    summaries of relevant code files. Attempt to answer the question and only respond if you're confident in the answer. \n",
            "    If you need any more information, respond with ONLY the phrase \"need more context\". \n",
            "\n",
            "    The user query is: What tracing frameworks are used across the codebase? \n",
            "\n",
            "    The summaries are: {'onlineboutique-codefiles/src/emailservice/requirements.in': 'This Python code uses the `google-cloud-profiler` and `google-cloud-trace` libraries to profile and trace Google Cloud applications. It also uses the `opentelemetry-distro` library for distributed tracing and the `grpcio` library for handling gRPC requests. The `requests` library is employed for making HTTP requests.', 'onlineboutique-codefiles/src/recommendationservice/requirements.in': 'This Python code uses the Google Cloud Profiler library to perform profiling and monitoring operations. It utilizes the gRPC health checking library for service health checks. The code also employs the OpenTelemetry instrumentation library for gRPC to enable tracing and observability.', 'onlineboutique-codefiles/src/checkoutservice/go.mod': 'This Go code is part of a microservice demonstration project. It uses the Cloud Profiler library for profiling, gRPC and gRPC-gateway for communication, OpenTelemetry for tracing and metrics, and Logrus for logging. The code includes functions such as `createShipment` and `getShipmentStatus` for managing checkouts.'}\n",
            "    \n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "'🐝 Completed at small, Gemini had enough context to respond. RESPONSE: \\nThe codebase uses `google-cloud-profiler`, `google-cloud-trace`, `opentelemetry-distro`, `grpcio`, and `requests` for profiling and tracing in Python, and `Cloud Profiler`, `grpc`, `grpc-gateway`, `OpenTelemetry`, and `Logrus` for profiling, tracing, logging, and health checking in Go.'"
            ]
          },
          "execution_count": 61,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "small_to_big(\"What tracing frameworks are used across the codebase?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 62,
      "metadata": {
        "id": "9ad4358c726d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "    You are a codebase helper. You will be given a user's question about the codebase, along with \n",
            "    summaries of relevant code files. Attempt to answer the question and only respond if you're confident in the answer. \n",
            "    If you need any more information, respond with ONLY the phrase \"need more context\". \n",
            "\n",
            "    The user query is: Describe in detail exactly how the ListRecommendations function works. \n",
            "\n",
            "    The summaries are: {'onlineboutique-codefiles/src/recommendationservice/client.py': 'This is a Python script that uses gRPC to communicate with a RecommendationService. It imports necessary libraries and sets up a server stub. The script forms a request and makes a call to the server, then logs the response using a custom JSON logger. It allows users to specify a port number as an argument.', 'onlineboutique-codefiles/src/productcatalogservice/.dockerignore': 'This is a directory listing. It is not a source code file, so I cannot provide a language identification or summary.', 'onlineboutique-codefiles/src/adservice/src/main/java/hipstershop/AdService.java': 'This is a Java gRPC server for an ad service.\\nIt has a `getAds` function that takes a request containing context and returns a response with a list of ads.\\nIt uses a map of categories to ads to retrieve the ads and a random function to retrieve random ads if no category is specified.'}\n",
            "    \n",
            "Prompted for a specific filename, Gemini said: onlineboutique-codefiles/src/recommendationservice/recommendation_server.py\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "\"🦖 Completed at big. RESPONSE: \\nThe `ListRecommendations` function in the `recommendationservice` retrieves a list of recommended products from the product catalog. It receives a request object containing a list of product IDs, and returns a response object containing a list of recommended product IDs.\\n\\nThe function works as follows:\\n\\n1. It fetches a list of all products from the product catalog using the `ListProducts` method of the `product_catalog_stub` gRPC client.\\n\\n2. It filters out any products that are already included in the request's list of product IDs, to avoid recommending products that the user has already seen.\\n\\n3. It randomly selects a subset of the remaining products to return as recommendations.\\n\\n4. It builds and returns a response object containing the list of recommended product IDs.\""
            ]
          },
          "execution_count": 62,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "small_to_big(\"Describe in detail exactly how the ListRecommendations function works.\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "small_to_big_rag.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
