{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ijGzTHJJUCPY"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VEqbX8OhE8y9"
      },
      "source": [
        "# Building a Conversational Search Agent with Agent Engine and RAG on Vertex AI Search\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_vertex_ai_search_rag_agent.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/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fagent-engine%2Ftutorial_vertex_ai_search_rag_agent.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://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_vertex_ai_search_rag_agent.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",
        "  <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/agent-engine/tutorial_vertex_ai_search_rag_agent.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\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/agent-engine/tutorial_vertex_ai_search_rag_agent.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/agent-engine/tutorial_vertex_ai_search_rag_agent.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/agent-engine/tutorial_vertex_ai_search_rag_agent.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/agent-engine/tutorial_vertex_ai_search_rag_agent.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/agent-engine/tutorial_vertex_ai_search_rag_agent.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>            \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84e7e432e6ff"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Kristopher Overholt](https://github.com/koverholt) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CkHPv2myT2cx"
      },
      "source": [
        "## Overview\n",
        "\n",
        "### Agent Engine in Vertex AI\n",
        "\n",
        "[Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview) is a managed service in Vertex AI that helps you to build and deploy agent frameworks. It gives you the flexibility to choose how much reasoning you want to delegate to the LLM and how much you want to handle with customized code. You can define Python functions that get used as tools via [Gemini Function Calling](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling). Agent Engine integrates closely with the Python SDK for the Gemini model in Vertex AI, and it can manage prompts, agents, and examples in a modular way. Agent Engine is compatible with LangChain, LlamaIndex, or other Python frameworks."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DrkcqHrrwMAo"
      },
      "source": [
        "### Objectives\n",
        "\n",
        "In this tutorial, you will build and deploy an agent (model, tools, and reasoning) using the Vertex AI SDK for Python.\n",
        "\n",
        "Your agent will use [LangChain](https://python.langchain.com/docs/get_started/introduction) and [Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/enterprise-search-introduction) to retrieve structured data indexed from the [Movies Dataset on Kaggle](https://www.kaggle.com/datasets/rounakbanik/the-movies-dataset?select=movies_metadata.csv) using retrieval augmented generation (RAG).\n",
        "\n",
        "- Install the Vertex AI SDK for Python\n",
        "- Define a model for your agent\n",
        "- Define Python functions as tools so that our agent can:\n",
        "   - Search and retrieve movie information from Vertex AI Search\n",
        "- Use the LangChain agent template provided in the Vertex AI SDK for Agent Engine\n",
        "- Test your agent locally before deploying\n",
        "- Deploy and test your agent on Agent Engine in Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C9nEPojogw-g"
      },
      "source": [
        "### Enable APIs and Services\n",
        "\n",
        "This tutorial uses the following billable components of Google Cloud, which you'll need to enable for this tutorial:\n",
        "\n",
        "- [Enable Vertex AI API](https://console.cloud.google.com/apis/api/aiplatform.googleapis.com/overview)\n",
        "- [Enable Vertex AI Search API](https://console.cloud.google.com/apis/api/discoveryengine.googleapis.com/overview)\n",
        "- [Enable Resource Manager API](https://console.cloud.google.com/apis/api/cloudresourcemanager.googleapis.com/overview)\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r11Gu7qNgx1p"
      },
      "source": [
        "## Getting Started\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK for Python\n",
        "\n",
        "Install the latest version of the Vertex AI SDK for Python and extra dependencies related to Agent Engine, LangChain, and Vertex AI Search:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet \\\n",
        "    \"google-cloud-aiplatform[agent_engines,langchain]\" \\\n",
        "    langchain-google-community \\\n",
        "    google-cloud-discoveryengine \\\n",
        "    google-api-python-client"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, run the following cell to authenticate your environment. This step is not required if you are using [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "STAGING_BUCKET = f\"gs://{PROJECT_ID}-agent-engine-staging\"  # @param {type:\"string\"}\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15e71c08d80d"
      },
      "source": [
        "## Example: Building and Deploying a Conversational Search Agent for Movies\n",
        "\n",
        "In this tutorial, you'll build Python functions as tools that the Gemini model will use to search for information about movies and ground its responses using a RAG implementation with Vertex AI Search."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jXHfaVS66_01"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "lslYAvw37JGQ"
      },
      "outputs": [],
      "source": [
        "import warnings\n",
        "\n",
        "from IPython.display import Markdown, display\n",
        "from googleapiclient import discovery\n",
        "from langchain.agents.format_scratchpad.tools import format_to_tool_messages\n",
        "from langchain.memory import ChatMessageHistory\n",
        "from langchain_core import prompts\n",
        "from vertexai.agent_engines import LangchainAgent\n",
        "\n",
        "warnings.filterwarnings(\"ignore\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "275640d559e9"
      },
      "source": [
        "### Define generative model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "47f07c83a754"
      },
      "source": [
        "The first step of your agent involves the generative model you want to use. Here you'll define the Gemini model for your agent:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "921890fcb875"
      },
      "outputs": [],
      "source": [
        "model = \"gemini-2.5-flash\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "89b4bac5179a"
      },
      "source": [
        "### Create a data store in Vertex AI Search\n",
        "\n",
        "Now you'll create a data store within Vertex AI Search and index records from a movie data set.\n",
        "\n",
        "Follow the tutorial steps in the Vertex AI Search documentation to [create a data store with structured data](https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#structured-data), then [create a search app](https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#structured-data_1) that points to that data store.\n",
        "\n",
        "You should also enable the **Enterprise edition features** and **Generative Responses** options so that you can use the [LangChain retriever for Vertex AI Search](https://python.langchain.com/docs/integrations/retrievers/google_vertex_ai_search/) within your agent to search indexed records.\n",
        "\n",
        "Once the import is complete, you can navigate to your data store's **Data** page and obtain the values of your data store ID and region, which you can paste into the cell below:\n",
        "\n",
        "<img width=\"60%\" src=\"https://storage.googleapis.com/github-repo/generative-ai/gemini/agent-engine/images/vertex-ai-search-data-store-id.png\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "ed87f5b9d63e"
      },
      "outputs": [],
      "source": [
        "DATA_STORE_ID = \"[your-data-store-id]\"  # @param {type:\"string\"}\n",
        "LOCATION_ID = \"global\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ea02086bd8e8"
      },
      "source": [
        "### Define Python functions as tools\n",
        "\n",
        "The second component of your agent involves Python functions as tools, which will enable the Gemini model to interact with external systems, databases, document stores, and other APIs so that the model can get the most up-to-date information or take action with those systems.\n",
        "\n",
        "In this example, you'll define a function that sends a query to Vertex AI Search and returns relevant records from the data store that you created in the previous section:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "766227a0e6cc"
      },
      "outputs": [],
      "source": [
        "def search_kaggle_movies(query: str) -> str:\n",
        "    \"\"\"Search across records in the Kaggle Movies data set.\"\"\"\n",
        "    from langchain_google_community import VertexAISearchRetriever\n",
        "\n",
        "    retriever = VertexAISearchRetriever(\n",
        "        project_id=PROJECT_ID,\n",
        "        data_store_id=DATA_STORE_ID,\n",
        "        location_id=LOCATION_ID,\n",
        "        engine_data_type=1,\n",
        "        max_documents=10,\n",
        "    )\n",
        "\n",
        "    result = str(retriever.invoke(query))\n",
        "    return result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "709bf1f8e8b0"
      },
      "source": [
        "Now you can test your search function with sample input to ensure that it's working as expected:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4e700cbb3c7a"
      },
      "outputs": [],
      "source": [
        "search_kaggle_movies(\"space exploration\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "99b00032a7aa"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ Important notes:</b><br>\n",
        "<br>\n",
        "<b>If you get an empty response when running the previous cell:</b><br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;Check the status of your data store in Vertex AI Search and try running the cell again after indexing is complete.<br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;Once you've created your search app and data store, you can view the <b>Activity</b> tab within your data store's <b>Data</b> page to see the status of your data ingestion.<br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;For this dataset, the import and indexing process takes about 5 to 10 minutes.<br>\n",
        "<br>\n",
        "<b>If you get an error when running the previous cell:</b><br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;In order for this sample notebook to work with the Google Vertex AI Search retriever in LangChain,<br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;you'll need to create a <a href=\"https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#create_a_data_store\">data store</a> <b>and</b> a <a href=\"https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#create_a_search_app\">search app</a> associated with it in Vertex AI Search.<br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;If you only create a data store, the LangChain retriever might throw errors when making queries against the data store.\n",
        "</div>\n",
        "\n",
        "</div>\n",
        "\n",
        "### Define agent\n",
        "\n",
        "The third component of your agent involves adding a reasoning layer, which helps your agent use the tools that you provided to help the end user achieve a higher-level goal.\n",
        "\n",
        "If you were to use Gemini and Function Calling on their own without a reasoning layer, you would need to handle the process of calling functions and APIs in your application code, and you would need to implement retries and additional logic to ensure that your function calling code is resilient to failures and malformed requests.\n",
        "\n",
        "Define the prompt template and initialize the chat session history:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "e1585c9c17de"
      },
      "outputs": [],
      "source": [
        "# Define prompt template\n",
        "prompt = {\n",
        "    \"history\": lambda x: x[\"history\"],\n",
        "    \"input\": lambda x: x[\"input\"],\n",
        "    \"agent_scratchpad\": (lambda x: format_to_tool_messages(x[\"intermediate_steps\"])),\n",
        "} | prompts.ChatPromptTemplate.from_messages(\n",
        "    [\n",
        "        prompts.MessagesPlaceholder(variable_name=\"history\"),\n",
        "        (\"user\", \"{input}\"),\n",
        "        prompts.MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
        "    ]\n",
        ")\n",
        "\n",
        "# Initialize session history\n",
        "store = {}\n",
        "\n",
        "\n",
        "def get_session_history(session_id: str):\n",
        "    if session_id not in store:\n",
        "        store[session_id] = ChatMessageHistory()\n",
        "    return store[session_id]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8e2f13d9a7df"
      },
      "source": [
        "Now you'll use the LangChain agent template provided in the Vertex AI SDK for Agent Engine, which brings together the model, tools, and reasoning that you've built up so far:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "7fe4fa1e163d"
      },
      "outputs": [],
      "source": [
        "agent = LangchainAgent(\n",
        "    prompt=prompt,\n",
        "    model=model,\n",
        "    chat_history=get_session_history,\n",
        "    model_kwargs={\"temperature\": 0},\n",
        "    tools=[search_kaggle_movies],\n",
        "    agent_executor_kwargs={\"return_intermediate_steps\": True},\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "96a94352da0f"
      },
      "source": [
        "### Test your agent locally\n",
        "\n",
        "Now you can test the model and agent behavior to ensure that it's working as expected before you deploy it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "03b16998bc2e"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": "Here are some movies featuring robots:\n\n*   **Robots** (2005): Rodney Copperbottom is a young robot inventor who dreams of making the world a better place, until the evil Ratchet takes over Big Weld Industries. Now, Rodney's dreams – and those of his friends – are in danger of becoming obsolete.\n*   **The War of the Robots** (1978): An alien civilization, which facing eminent extinction, kidnaps two famous genetic scientists from Earth. A troop of soldiers is dispatched to combat the humanoid robots and rescue the victims.\n*   **Robot Stories** (2003): Four stories including: \"My Robot Baby,\" in which a couple must care for a robot baby before adopting a human child; \"The Robot Fixer,\" in which a mother tries to connect with her dying son; \"Machine Love,\" in which an office worker android learns that he, too, needs love; and \"Clay,\" in which an old sculptor must choose between natural death and digital immortality.\n*   **I, Robot** (2004): In 2035, where robots are common-place and abide by the three laws of robotics, a techno-phobic cop investigates an apparent suicide. Suspecting that a robot may be responsible for the death, his investigation leads him to believe that humanity may be in danger.\n*   **Robot Wars** (1993): In the year 2041, the rebel Centros are a plague to the survivors of the great toxic gas scare of 1993. A renegade Megarobot pilot and an archaeologist must team up (despite personal differences, a reluctant romance, and official pressure to cease and desist) to thwart the Centro's attempts to resurrect a hidden Megarobot, with which they can challenge the prevailing order.\n*   **Robot Monster** (1953): Ro-Man, an alien robot who greatly resembles a gorilla in a diving helmet, is sent to earth to destroy all human life. Ro-Man falls in love with one of the last six remaining humans, and struggles to understand how his programming can instruct him to kill her while his heart demands that he can't.\n*   **Robot Overlords** (2015): Earth has been conquered by robots from another galaxy and the human survivors must stay in their homes, or risk incineration.\n*   **OMG, I'm a Robot!** (2015): A sensitive guy finds out he's... a robot.\n*   **Robot & Frank** (2012): A delightful dramatic comedy, a buddy picture, and, for good measure, a heist film. Curmudgeonly old Frank lives by himself. His routine involves daily visits to his local library, where he has a twinkle in his eye for the librarian. His grown children are concerned about their father's well-being and buy him a caretaker robot. Initially resistant to the idea, Frank soon appreciates the benefits of robotic support – like nutritious meals and a clean house – and eventually begins to treat his robot like a true companion. With his robot's assistance, Frank's passion for his old, unlawful profession is reignited, for better or worse.\n*   **Robot Carnival** (1987): An anthology of various tales told in various styles with robots being the one common element among them.",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = agent.query(\n",
        "    input=\"Tell me about movies featuring robots\",\n",
        "    config={\"configurable\": {\"session_id\": \"demo\"}},\n",
        ")\n",
        "\n",
        "display(Markdown(response[\"output\"]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9abc4a214f3c"
      },
      "source": [
        "### Deploy your agent on Vertex AI\n",
        "\n",
        "Now that you've tested your agent locally, you're ready to deploy it to Agent Engine in Vertex AI. This will make your agent accessible remotely and allow you to integrate it into larger systems or provide it as a service."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c3325266144e"
      },
      "outputs": [],
      "source": [
        "client = vertexai.Client(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "remote_agent = client.agent_engines.create(\n",
        "    agent=agent,\n",
        "    config={\n",
        "        \"staging_bucket\": STAGING_BUCKET,\n",
        "        \"requirements\": [\n",
        "            \"google-cloud-aiplatform[agent_engines,langchain]\",\n",
        "            \"langchain-google-community\",\n",
        "            \"google-cloud-discoveryengine\",\n",
        "        ],\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "db896f74fa3b"
      },
      "source": [
        "### Grant Discovery Engine Editor access to Agent Engine service account\n",
        "\n",
        "Before you send queries to your remote agent, you'll need to grant the **Discovery Engine Editor** role to the Agent Engine service account.\n",
        "\n",
        "After you've completed this step, you remote agent will be able to retrieve documents from the data store that you created in Vertex AI Search:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "75ac983656be"
      },
      "outputs": [],
      "source": [
        "# Retrieve the project number associated with your project ID\n",
        "service = discovery.build(\"cloudresourcemanager\", \"v1\")\n",
        "request = service.projects().get(projectId=PROJECT_ID)\n",
        "response = request.execute()\n",
        "project_number = response[\"projectNumber\"]\n",
        "project_number"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "53b3f4665183"
      },
      "outputs": [],
      "source": [
        "# Add a new role binding to the IAM policy\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:service-{project_number}@gcp-sa-aiplatform-re.iam.gserviceaccount.com \\\n",
        "    --role=roles/discoveryengine.editor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "efee31ce8859"
      },
      "source": [
        "### Test your remotely deployed agent\n",
        "\n",
        "With your conversational search agent deployed, you can send prompts to test that it's working as expected and that it can retrieve movie data from your Vertex AI Search data store:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "5f08c5685b4e"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": "Here are some movies featuring robots:\n\n*   **Robots** (2005): Rodney Copperbottom is a young robot inventor who dreams of making the world a better place, until the evil Ratchet takes over Big Weld Industries. Now, Rodney's dreams – and those of his friends – are in danger of becoming obsolete.\n*   **The War of the Robots** (1978): An alien civilization, which facing eminent extinction, kidnaps two famous genetic scientists from Earth. A troop of soldiers is dispatched to combat the humanoid robots and rescue the victims.\n*   **Robot Stories** (2003): Four stories including: \"My Robot Baby,\" in which a couple must care for a robot baby before adopting a human child; \"The Robot Fixer,\" in which a mother tries to connect with her dying son; \"Machine Love,\" in which an office worker android learns that he, too, needs love; and \"Clay,\" in which an old sculptor must choose between natural death and digital immortality.\n*   **I, Robot** (2004): In 2035, where robots are common-place and abide by the three laws of robotics, a techno-phobic cop investigates an apparent suicide. Suspecting that a robot may be responsible for the death, his investigation leads him to believe that humanity may be in danger.\n*   **Robot Wars** (1993): In the year 2041, the rebel Centros are a plague to the survivors of the great toxic gas scare of 1993. A renegade Megarobot pilot and an archaeologist must team up (despite personal differences, a reluctant romance, and official pressure to cease and desist) to thwart the Centro's attempts to resurrect a hidden Megarobot, with which they can challenge the prevailing order.\n*   **Robot Monster** (1953): Ro-Man, an alien robot who greatly resembles a gorilla in a diving helmet, is sent to earth to destroy all human life. Ro-Man falls in love with one of the last six remaining humans, and struggles to understand how his programming can instruct him to kill her while his heart demands that he can't.\n*   **Robot Overlords** (2015): Earth has been conquered by robots from another galaxy and the human survivors must stay in their homes, or risk incineration.\n*   **OMG, I'm a Robot!** (2015): A sensitive guy finds out he's... a robot.\n*   **Robot & Frank** (2012): A delightful dramatic comedy, a buddy picture, and, for good measure, a heist film. Curmudgeonly old Frank lives by himself. His routine involves daily visits to his local library, where he has a twinkle in his eye for the librarian. His grown children are concerned about their father's well-being and buy him a caretaker robot. Initially resistant to the idea, Frank soon appreciates the benefits of robotic support – like nutritious meals and a clean house – and eventually begins to treat his robot like a true companion. With his robot's assistance, Frank's passion for his old, unlawful profession is reignited, for better or worse.\n*   **Robot Carnival** (1987): An anthology of various tales told in various styles with robots being the one common element among them.\n*   **Robot Holocaust** (1986): Neo, a drifter from the atomic-blasted wastelands, arrives with his klutzy robot sidekick at a factory where slaves labor to fuel the sinister Dark One's Power Station. There, he meets a comely woman who convinces him to help rescue her scientist father, who has invented a device that can break the Dark One's control over the slaves. Gathering a motley crew of allies on the way, Neo and pals travel to the Power Station, where they confront the Dark One's evil servants.\n*   **Robotropolis** (2011): A group of reporters are covering the unveiling of a new facility that is completely maintained by robot prototypes. When one of the robots goes haywire, the reporters find themselves not just reporting on the malfunction, but fighting for their lives.\n*   **They Call Me Jeeg Robot** (2015): Enzo Ceccotti comes into contact with a radioactive substance, then accidently discovers he has superpowers. A touchy, navel-gazing introvert, he's sure his new capabilities will do wonders for his life of crime, but that all changes when he meets Alessia, who's convinced he's the hero from the famous Japanese comic strip, Steel Jeeg Robot.\n*   **Wild, Wild Planet** (1965): A rogue cop must stop a scientist from taking over the world with his deadly female robots, who are shrinking the world leaders.",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = remote_agent.query(\n",
        "    input=\"Tell me about movies featuring robots\",\n",
        "    config={\"configurable\": {\"session_id\": \"demo\"}},\n",
        ")\n",
        "\n",
        "display(Markdown(response[\"output\"]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "da6076427022"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": "\"I, Robot\" was released on July 15, 2004. It is an action and science fiction film with a runtime of 115 minutes. The movie's tagline is \"Laws are made to be broken.\"\n\nThe plot is set in 2035, where robots are commonplace and abide by the three laws of robotics. A techno-phobic cop investigates an apparent suicide, and suspecting that a robot may be responsible for the death, his investigation leads him to believe that humanity may be in danger.\n\nThe film had a budget of $120,000,000 and grossed $347,234,916. It was produced by Twentieth Century Fox Film Corporation, Laurence Mark Productions, Davis Entertainment, Overbrook Entertainment, Mediastream Vierte Film GmbH & Co. Vermarktungs KG, and Canlaws Productions. The primary spoken language is English, and it was produced in the United States of America and Germany. The movie has a vote average of 6.7 based on 3889 votes.",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = remote_agent.query(\n",
        "    input=\"Tell me more about the movie I, Robot\",\n",
        "    config={\"configurable\": {\"session_id\": \"demo\"}},\n",
        ")\n",
        "\n",
        "display(Markdown(response[\"output\"]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "797098dd1f88"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": "Based on the descriptions, here are the robot movies that appear to be comedies:\n\n*   **Robots** (2005): This is an animated film, often implying a comedic tone, and the description suggests a lighthearted adventure.\n*   **OMG, I'm a Robot!** (2015): The title itself suggests a comedic premise.\n*   **Robot & Frank** (2012): The description explicitly states it's \"A delightful dramatic comedy.\"\n*   **Robot Monster** (1953): While a sci-fi film, its premise of a gorilla-like robot falling in love and struggling with programming vs. heart often leads to it being viewed as a cult classic with unintentional comedic elements due to its B-movie nature.\n*   **Robot Carnival** (1987): As an anthology, it's possible some segments are comedic, but the overall description doesn't explicitly state it. However, animated anthologies often include comedic shorts.\n*   **They Call Me Jeeg Robot** (2015): While it has action, the premise of a \"touchy, navel-gazing introvert\" who gets superpowers and is mistaken for a comic book hero suggests comedic elements.\n\nThe most explicitly comedic ones are **Robots**, **OMG, I'm a Robot!**, and **Robot & Frank**.",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = remote_agent.query(\n",
        "    input=\"Which of those robot movies are comedies?\",\n",
        "    config={\"configurable\": {\"session_id\": \"demo\"}},\n",
        ")\n",
        "\n",
        "display(Markdown(response[\"output\"]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9f2f7d3ed7bd"
      },
      "source": [
        "### Querying your deployed agent\n",
        "\n",
        "You've now deployed your Agent Engine agent and can [interact with it in multiple ways](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/use/overview), both within this notebook and from other applications or environments. The primary methods for accessing your deployed agent are via the Python client library or through REST API calls. Here's an overview of both methods:\n",
        "\n",
        "**Method 1: Reusing within this notebook or another Python environment**\n",
        "\n",
        "You can directly reuse and query the `remote_agent` instance you created in this notebook.\n",
        "\n",
        "Or, you can instantiate a new instance in another notebook or Python script. To do this, you'll need to retrieve your deployed agent's resource name that uniquely identifies your agent, which is a string that includes the project, location, and Agent Engine ID. You can retrieve it by running the following code in the notebook or environment where you created your agent:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fdaf8b91413f"
      },
      "outputs": [],
      "source": [
        "remote_agent.api_resource.name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "060f8369d113"
      },
      "source": [
        "Use the resource name to load the agent in your other notebook or Python script, then query the remote agent as usual:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "78af4442827e"
      },
      "outputs": [],
      "source": [
        "# from vertexai import agent_engines\n",
        "\n",
        "# client = vertexai.Client(project=PROJECT_ID, location=LOCATION)\n",
        "# AGENT_ENGINE_RESOURCE_NAME = \"YOUR_AGENT_ENGINE_RESOURCE_NAME\"  # Replace with the resource name of your deployed Agent Engine\n",
        "\n",
        "# remote_agent = client.agent_engines.get(name=AGENT_ENGINE_RESOURCE_NAME)\n",
        "# response = remote_agent.query(input=\"List some sci-fi movies from the 1990s\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "63ab06554fc0"
      },
      "source": [
        "**Method 2: Accessing from other environments via REST API**\n",
        "\n",
        "Beyond the Python client library, your deployed Vertex AI agent can be [queried using REST API calls](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/use/overview), including:\n",
        "\n",
        "- Python: You can use Python's `requests` library or similar tools to make HTTP calls to the Vertex AI REST API.\n",
        "- cURL: A command-line tool, cURL allows you to send HTTP requests directly. This is useful for testing and debugging.\n",
        "- Other Programming Languages: If you prefer a different language for your application, you can use its native HTTP client library to make REST API calls.\n",
        "\n",
        "In summary, you have access to your deployed Agent Engine agent through the Python client library within Python environments, and more universally through its REST API via tools and programming languages of your choosing."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "806f4e063a52"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "After you've finished experimenting, it's a good practice to clean up your cloud resources. You can delete the deployed Agent Engine instance and optionally remove the staging bucket to avoid any unexpected charges on your Google Cloud account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7cec827288b7"
      },
      "outputs": [],
      "source": [
        "# Delete the deployed agent\n",
        "client.agent_engines.delete(name=remote_agent.api_resource.name)\n",
        "\n",
        "# Optionally, delete the staging bucket\n",
        "# from google.cloud import storage\n",
        "# storage.Client().bucket(STAGING_BUCKET.replace(\"gs://\", \"\")).delete(force=True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tutorial_vertex_ai_search_rag_agent.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
