gemini/reasoning-engine/tutorial_vertex_ai_search_rag_agent.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "ijGzTHJJUCPY",
    "tags": []
   },
   "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 Reasoning 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/reasoning-engine/tutorial_vertex_ai_search_rag_agent.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" 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%2Freasoning-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/reasoning-engine/tutorial_vertex_ai_search_rag_agent.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" 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/reasoning-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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| | |\n",
    "|-|-|\n",
    "|Author(s) | [Kristopher Overholt](https://github.com/koverholt) |"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "CkHPv2myT2cx"
   },
   "source": [
    "## Overview\n",
    "\n",
    "### Reasoning Engine in Vertex AI\n",
    "\n",
    "[Reasoning Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/overview) (LangChain on Vertex AI) is a managed service in Vertex AI that helps you to build and deploy an agent reasoning framework. 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). Reasoning 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. Reasoning Engine is compatible with LangChain, LlamaIndex, or other Python frameworks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DrkcqHrrwMAo",
    "jp-MarkdownHeadingCollapsed": true
   },
   "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 Reasoning Engine\n",
    "- Test your agent locally before deploying\n",
    "- Deploy and test your agent on Reasoning 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 Reasoning Engine, LangChain, and Vertex AI Search:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "tFy3H3aPgx12",
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade --quiet \\\n",
    "    \"google-cloud-aiplatform[langchain,reasoningengine]\" \\\n",
    "    cloudpickle==3.0.0 \\\n",
    "    pydantic==2.7.4 \\\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": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XRvKdaPDTznN",
    "outputId": "154a71b5-f302-4f53-ed2f-b3e5fef9195b",
    "tags": []
   },
   "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",
    "\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",
    "tags": []
   },
   "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",
    "tags": []
   },
   "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": {},
   "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",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from googleapiclient import discovery\n",
    "from IPython.display import display, Markdown\n",
    "from langchain.agents.format_scratchpad.tools import format_to_tool_messages\n",
    "from langchain_core import prompts\n",
    "from langchain.memory import ChatMessageHistory\n",
    "from vertexai.preview import reasoning_engines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define generative model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 1.5 Pro model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = \"gemini-1.5-pro-001\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 and preview a search app for structured data](https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#create_and_preview_a_search_app_for_structured_data_from). You can use a company name of **Movie Information** in your search app. 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 you've created your search app and data store, you can view the **Activity** tab within your data store's **Data** page to see the status of your data ingestion. For this dataset, the import and indexing process takes about ten minutes.\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/reasoning-engine/images/vertex-ai-search-data-store-id.png\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_STORE_ID = \"[your-data-store-id]\"  # @param {type:\"string\"}\n",
    "LOCATION_ID = \"global\"  # @param {type:\"string\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": {},
   "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": {},
   "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": {},
   "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": {},
   "source": [
    "### 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": {},
   "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": {},
   "source": [
    "Now you'll use the LangChain agent template provided in the Vertex AI SDK for Reasoning Engine, which brings together the model, tools, and reasoning that you've built up so far:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = reasoning_engines.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": {},
   "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": {},
   "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": {},
   "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": {},
   "outputs": [],
   "source": [
    "agent = reasoning_engines.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": {},
   "source": [
    "Now you're ready to deploy your agent to Reasoning Engine in Vertex AI by calling `reasoning_engines.ReasoningEngine.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": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using bucket your-bucket-name\n",
      "Writing to gs://your-bucket-name/reasoning_engine/reasoning_engine.pkl\n",
      "Writing to gs://your-bucket-name/reasoning_engine/requirements.txt\n",
      "Creating in-memory tarfile of extra_packages\n",
      "Writing to gs://your-bucket-name/reasoning_engine/dependencies.tar.gz\n",
      "Creating ReasoningEngine\n",
      "Create ReasoningEngine backing LRO: projects/954731510984/locations/us-central1/reasoningEngines/5997035485053059072/operations/2550512418300100608\n",
      "ReasoningEngine created. Resource name: projects/954731510984/locations/us-central1/reasoningEngines/5997035485053059072\n",
      "To use this ReasoningEngine in another session:\n",
      "reasoning_engine = vertexai.preview.reasoning_engines.ReasoningEngine('projects/954731510984/locations/us-central1/reasoningEngines/5997035485053059072')\n"
     ]
    }
   ],
   "source": [
    "remote_agent = reasoning_engines.ReasoningEngine.create(\n",
    "    agent,\n",
    "    requirements=[\n",
    "        \"google-cloud-aiplatform[langchain,reasoningengine]\",\n",
    "        \"cloudpickle==3.0.0\",\n",
    "        \"pydantic==2.7.4\",\n",
    "        \"langchain-google-community\",\n",
    "        \"google-cloud-discoveryengine\",\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Grant Discovery Engine Editor access to Reasoning 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 Reasoning 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": {},
   "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": {},
   "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": {},
   "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": {},
   "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": {},
   "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": {},
   "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": {},
   "source": [
    "### Reusing your deployed agent from other applications or SDKs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now import and use the remotely deployed Reasoning Engine in this notebook session or in a different notebook or Python script by uncommenting and adapting the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from vertexai.preview import reasoning_engines\n",
    "\n",
    "# PROJECT_ID = \"YOUR_PROJECT_ID\"\n",
    "# LOCATION = \"YOUR_LOCATION\"\n",
    "# REASONING_ENGINE_ID = \"YOUR_REASONING_ENGINE_ID\"\n",
    "\n",
    "# remote_agent = reasoning_engines.ReasoningEngine(f\"projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{REASONING_ENGINE_ID}\")\n",
    "# response = remote_agent.query(input=query)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or, you can query your agent from other programming languages using any of the [available client libraries in Vertex AI](https://cloud.google.com/vertex-ai/docs/start/client-libraries), including C#, Java, Node.js, Python, Go, or REST API."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m115",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m115"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}