{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "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) |"
      ]
    },
    {
      "attachments": {},
      "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": 2,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet \\\n",
        "    \"google-cloud-aiplatform[agent_engines,langchain]\" \\\n",
        "    cloudpickle==3.0.0 \\\n",
        "    \"pydantic>=2.10\" \\\n",
        "    langchain-google-community \\\n",
        "    google-cloud-discoveryengine \\\n",
        "    google-api-python-client"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart current 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 will restart the current kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "# Restart kernel after installs so that your environment can access the new packages\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you 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": 4,
      "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": 5,
      "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 = \"gs://[your-staging-bucket]\"  # @param {type:\"string\"}\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "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": 6,
      "metadata": {
        "id": "lslYAvw37JGQ"
      },
      "outputs": [],
      "source": [
        "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 import agent_engines\n",
        "from vertexai.preview.reasoning_engines import LangchainAgent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "275640d559e9"
      },
      "source": [
        "### Define generative model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "47f07c83a754"
      },
      "source": [
        "The first component of your agent involves the version of the generative model you want to use in your agent. Here you'll use the Gemini 2.0 model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "921890fcb875"
      },
      "outputs": [],
      "source": [
        "model = \"gemini-2.0-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 **Advanced LLM features** 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 in 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": 8,
      "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": 9,
      "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": 10,
      "metadata": {
        "id": "4e700cbb3c7a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'[Document(page_content=\\'{\"runtime\": \"152\", \"tagline\": \"Let the Magic Begin.\", \"poster_path\": \"https://image.tmdb.org/t/p/original/wuMc08IPKEatf9rnMNXvIDxqP4W.jpg\", \"id\": \"671\", \"budget\": \"125000000\", \"original_language\": \"en\", \"status\": \"Released\", \"homepage\": \"http://harrypotter.warnerbros.com/harrypotterandthedeathlyhallows/mainsite/index.html\", \"popularity\": 38.187238, \"belongs_to_collection\": \"{\\\\\\'id\\\\\\': 1241, \\\\\\'name\\\\\\': \\\\\\'Harry Potter Collection\\\\\\', \\\\\\'poster_path\\\\\\': \\\\\\'/8Xmkc1HvCOpMlbFvVabrtr6HAsp.jpg\\\\\\', \\\\\\'backdrop_path\\\\\\': \\\\\\'/wfnMt6LGqYHcNyOfsuusw5lX3bL.jpg\\\\\\'}\", \"original_title\": \"Harry Potter and the Philosopher\\\\\\'s Stone\", \"vote_count\": \"7188\", \"imdb_id\": \"tt0241527\", \"release_date\": \"2001-11-16\", \"spoken_languages\": \"[{\\\\\\'iso_639_1\\\\\\': \\\\\\'en\\\\\\', \\\\\\'name\\\\\\': \\\\\\'English\\\\\\'}]\", \"production_companies\": \"[{\\\\\\'name\\\\\\': \\\\\\'1492 Pictures\\\\\\', \\\\\\'id\\\\\\': 436}, {\\\\\\'name\\\\\\': \\\\\\'Warner Bros.\\\\\\', \\\\\\'id\\\\\\': 6194}, {\\\\\\'name\\\\\\': \\\\\\'Heyday Films\\\\\\', \\\\\\'id\\\\\\': 7364}]\", \"overview\": \"Harry Potter has lived under the stairs at his aunt and uncle\\\\\\'s house his whole life. But on his 11th birthday, he learns he\\\\\\'s a powerful wizard -- with a place waiting for him at the Hogwarts School of Witchcraft and Wizardry. As he learns to harness his newfound powers with the help of the school\\\\\\'s kindly headmaster, Harry uncovers the truth about his parents\\\\\\' deaths -- and about the villain who\\\\\\'s to blame.\", [...]'"
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "search_kaggle_movies(\"Harry Potter\")"
      ]
    },
    {
      "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": 11,
      "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": 12,
      "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": 13,
      "metadata": {
        "id": "03b16998bc2e"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "Here are some sci-fi movies from the 1990s:\n",
              "\n",
              "* **The Matrix (1999):** A computer hacker learns from mysterious rebels about the true nature of his reality and his role in the war against its controllers.\n",
              "* **Star Wars: Episode I - The Phantom Menace (1999):** Two Jedi Knights escape a hostile blockade to find a queen and her protector, and come across a young boy who may bring balance to the Force, but the long-dormant Sith resurface to claim their original power. \n",
              "* **Men in Black (1997):** A police officer joins a secret organization that polices and monitors extraterrestrial interactions on Earth. \n",
              "* **Starship Troopers (1997):** Humans in a fascist militaristic future wage war with giant alien bugs.\n",
              "* **Gatace (1997):** A genetically inferior man assumes the identity of a superior one in order to pursue his lifelong dream of space travel. \n",
              "* **12 Monkeys (1995):** In the year 2035, a convict is sent back in time to gather information about the man-made virus that wiped out most of the human"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = agent.query(\n",
        "    input=\"List some sci-fi movies from the 1990s\",\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 specified a model, tools, and reasoning for your agent and tested it out, you're ready to deploy your agent as a remote service in Vertex AI!\n",
        "\n",
        "You can re-define the agent to avoid any stateful information in the agent due to our testing in the previous cell:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "ae9d1c29fdd4"
      },
      "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": "e14c4fe23fa3"
      },
      "source": [
        "Now you're ready to deploy your agent to Agent Engine in Vertex AI by calling `agent_engines.create()` along with the instance of your agent and the Python packages that your agent requires at runtime:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "c3325266144e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using bucket your-bucket-name\n",
            "Writing to gs://your-bucket-name/agent_engine/agent_engine.pkl\n",
            "Writing to gs://your-bucket-name/agent_engine/requirements.txt\n",
            "Creating in-memory tarfile of extra_packages\n",
            "Writing to gs://your-bucket-name/agent_engine/dependencies.tar.gz\n",
            "Creating AgentEngine\n",
            "Create AgentEngine backing LRO: projects/your-project-number/locations/us-central1/reasoningEngines/5997035485053059072/operations/2550512418300100608\n",
            "AgentEngine created. Resource name: projects/your-project-number/locations/us-central1/reasoningEngines/5997035485053059072\n",
            "To use this AgentEngine in another session:\n",
            "agent_engine = vertexai.agent_engines.get('projects/your-project-number/locations/us-central1/reasoningEngines/5997035485053059072')\n"
          ]
        }
      ],
      "source": [
        "remote_agent = agent_engines.create(\n",
        "    agent,\n",
        "    requirements=[\n",
        "        \"google-cloud-aiplatform[agent_engines,langchain]\",\n",
        "        \"cloudpickle==3.0.0\",\n",
        "        \"pydantic>=2.10\",\n",
        "        \"langchain-google-community\",\n",
        "        \"google-cloud-discoveryengine\",\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": 16,
      "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": 17,
      "metadata": {
        "id": "53b3f4665183"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Updated IAM policy for project [your-project-id].\n",
            "auditConfigs:\n",
            "- auditLogConfigs:\n",
            "  - logType: ADMIN_READ\n",
            "  - logType: DATA_READ\n",
            "  - logType: DATA_WRITE\n",
            "  service: allServices\n",
            "bindings:\n",
            "- members:\n",
            "  - serviceAccount:service-954731510984@gcp-sa-aiplatform-re.iam.gserviceaccount.com\n",
            "  role: roles/discoveryengine.editor\n",
            "etag: BwYX4TeUgNo=\n",
            "version: 1\n"
          ]
        }
      ],
      "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 all of the core components of your community solar planning agent in place, you can send prompts to your remotely deployed agent to perform different tasks and test that it's working as expected:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "5f08c5685b4e"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "Here are some sci-fi movies from the 1990s:\n",
              "\n",
              "* **The Matrix (1999):** A computer hacker learns from mysterious rebels about the true nature of his reality and his role in the war against its controllers.\n",
              "* **Star Wars: Episode I - The Phantom Menace (1999):** Two Jedi Knights escape a hostile blockade to find a queen and her protector, and come across a young boy who may bring balance to the Force, but the long-dormant Sith resurface to claim their original power. \n",
              "* **Men in Black (1997):** A police officer joins a secret organization that polices and monitors extraterrestrial interactions on Earth. \n",
              "* **Starship Troopers (1997):** Humans in a fascist militaristic future wage war with giant alien bugs.\n",
              "* **Gatace (1997):** A genetically inferior man assumes the identity of a superior one in order to pursue his lifelong dream of space travel. \n",
              "* **12 Monkeys (1995):** In the year 2035, a convict is sent back in time to gather information about the man-made virus that wiped out most of the human"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = remote_agent.query(\n",
        "    input=\"List some sci-fi movies from the 1990s\",\n",
        "    config={\"configurable\": {\"session_id\": \"demo\"}},\n",
        ")\n",
        "\n",
        "display(Markdown(response[\"output\"]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "da6076427022"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "The main actors in The Matrix (1999) are:\n",
              "\n",
              "* **Keanu Reeves** as Neo\n",
              "* **Laurence Fishburne** as Morpheus \n",
              "* **Carrie-Anne Moss** as Trinity\n",
              "* **Hugo Weaving** as Agent Smith \n",
              "* **Joe Pantoliano** as Cypher \n"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = remote_agent.query(\n",
        "    input=\"Who are the actors in The Matrix?\",\n",
        "    config={\"configurable\": {\"session_id\": \"demo\"}},\n",
        ")\n",
        "\n",
        "display(Markdown(response[\"output\"]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "797098dd1f88"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "Yes, all of those actors have been in many other movies! Here are a few examples for each:\n",
              "\n",
              "* **Keanu Reeves:** John Wick (2014), The Devil's Advocate (1997), Speed (1994), Bill & Ted's Excellent Adventure (1989)\n",
              "* **Laurence Fishburne:**  Apocalypse Now (1979), Boyz n the Hood (1991), What's Love Got to Do with It (1993), Contagion (2011)\n",
              "* **Carrie-Anne Moss:** Memento (2000), Chocolat (2000), Disturbia (2007), Jessica Jones (TV series, 2015-2019)\n",
              "* **Hugo Weaving:** The Lord of the Rings trilogy (2001-2003), V for Vendetta (2005), Captain America: The First Avenger (2011), The Hobbit trilogy (2012-2014)\n",
              "* **Joe Pantoliano:** The Goonies (1985), Risky Business (1983), The Fugitive (1993), Memento (2000) \n"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = remote_agent.query(\n",
        "    input=\"Are those actors in any other movies?\",\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": 21,
      "metadata": {
        "id": "fdaf8b91413f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'projects/your-project-number/locations/us-central1/reasoningEngines/223106848019812424'"
            ]
          },
          "execution_count": 21,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "remote_agent.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": 22,
      "metadata": {
        "id": "78af4442827e"
      },
      "outputs": [],
      "source": [
        "# from vertexai.preview import agent_engines\n",
        "\n",
        "# AGENT_ENGINE_RESOURCE_NAME = \"YOUR_AGENT_ENGINE_RESOURCE_NAME\"  # Replace with the resource name of your deployed Agent Engine\n",
        "\n",
        "# remote_agent = agent_engines.get(AGENT_ENGINE_RESOURCE_NAME)\n",
        "# response = remote_agent.query(input=query)"
      ]
    },
    {
      "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 to avoid any unexpected charges on your Google Cloud account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "7cec827288b7"
      },
      "outputs": [],
      "source": [
        "remote_agent.delete()"
      ]
    }
  ],
  "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
}
