{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KWS7OXcEJptT"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 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": [
        "# Deploying an Agent with Agent Engine and MCP Toolbox for Databases\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_mcp_toolbox_for_databases.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fagent-engine%2Ftutorial_mcp_toolbox_for_databases.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/agent-engine/tutorial_mcp_toolbox_for_databases.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_mcp_toolbox_for_databases.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://goo.gle/4jeQzJW\">\n",
        "      <img width=\"32px\" src=\"https://cdn.qwiklabs.com/assets/gcp_cloud-e3a77215f0b8bfa9b3f611c0d2208c7e8708ed31.svg\" alt=\"Google Cloud logo\"><br> Open in  Cloud Skills Boost\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_mcp_toolbox_for_databases.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_mcp_toolbox_for_databases.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_mcp_toolbox_for_databases.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_mcp_toolbox_for_databases.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_mcp_toolbox_for_databases.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "23896bafa4c4"
      },
      "source": [
        "| Author |\n",
        "| --- |\n",
        "| [Shawn Yang](https://github.com/shawn-yang-google) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m62Ow6OrhFod"
      },
      "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 that helps you to build and deploy an agent 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). 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.\n",
        "\n",
        "\n",
        "## MCP Toolbox for Databases\n",
        "\n",
        "[MCP Toolbox for Databases](https://googleapis.github.io/genai-toolbox/getting-started/introduction/) is an open source MCP server for databases. It was designed with enterprise-grade and production-quality in mind. It enables you to develop tools easier, faster, and more securely by handling the complexities such as connection pooling, authentication, and more.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S0KQAd76m-X4"
      },
      "source": [
        "## Objectives\n",
        "\n",
        "In this tutorial, you will learn how to build and deploy an agent (model, tools, and reasoning) using the Vertex AI SDK and MCP Toolbox for Databases.\n",
        "\n",
        "You'll build and deploy an agent that uses the Gemini model, Python functions as tools, and LangGraph for orchestration.\n",
        "\n",
        "You will complete the following tasks:\n",
        "\n",
        "- Install the Vertex AI SDK for Python\n",
        "- Use the LangGraph to build components of a simple agent\n",
        "- Deploy a PostgreSQL database using Cloud SQL.\n",
        "- Deploy Toolbox for Databases using Cloud Run\n",
        "- Test your agent locally before deploying\n",
        "- Deploy and test your agent remotely"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C9nEPojogw-g"
      },
      "source": [
        "## Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "- `Vertex AI` (for Agent Engine)\n",
        "- `Cloud Run` (for MCP Toolbox for Databases)\n",
        "- `Cloud SQL` (for PostgreSQL DB)\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
        "[Cloud Run pricing](https://cloud.google.com/run/pricing), [Cloud SQL for PostgreSQL pricing](https://cloud.google.com/sql/docs/postgres/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4ijtb_v3M1fe"
      },
      "source": [
        "# Getting Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6RZhCO_aJd6J"
      },
      "source": [
        "## Install libraries\n",
        "\n",
        "Install necessary libraries for Vertex AI Agent Engine, LangChain, Toolbox, and Cloud SQL integration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f8cE6RysJdiP"
      },
      "outputs": [],
      "source": [
        "%pip install -q -U \\\n",
        "    \"google-cloud-aiplatform[agent_engines,langchain]\" \\\n",
        "    toolbox-langchain \\\n",
        "    \"langchain-google-cloud-sql-pg>=0.10.0\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EqEC4KxbMxaf"
      },
      "source": [
        "## Set up global parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XtiB5-LVVkv0"
      },
      "outputs": [],
      "source": [
        "# @title Required User Input:\n",
        "# Please update these values for your environment.\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"} Your Google Cloud Project ID.\n",
        "STAGING_BUCKET_NAME = \"[your-staging-bucket-name]\"  # @param {type:\"string\"} Name (only) of your GCS bucket for staging.\n",
        "# fmt: on\n",
        "PASSWORD = input(\n",
        "    \"Please provide a password to be used for the 'postgres' database user: \"\n",
        ")  # Securely input the database password."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7VClYJJuHw8P"
      },
      "outputs": [],
      "source": [
        "# @title Tutorial Defaults:\n",
        "# These parameters are pre-configured for this tutorial. Usually keep these unless you have specific needs.\n",
        "# fmt: off\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"} The region where resources will be created.\n",
        "INSTANCE = \"hotel-booking-instance\"  # @param {type:\"string\"} Name for the Cloud SQL instance.\n",
        "DATABASE = \"postgres\"  # @param {type:\"string\"} Name for the database within the instance.\n",
        "TABLE_NAME = \"hotels\"  # @param {type:\"string\"} Name for the table to be created.\n",
        "USER = \"postgres\"  # @param {type:\"string\"} Database admin username.\n",
        "TOOLBOX_SERVICE_ACCOUNT = \"toolbox-identity\"  # @param {type:\"string\"} Name for the toolbox service account.\n",
        "# fmt: on"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HHggtDkRID-j"
      },
      "outputs": [],
      "source": [
        "# @title Automatically Derived Parameters\n",
        "# These are constructed based on your input and defaults.\n",
        "STAGING_BUCKET = f\"gs://{STAGING_BUCKET_NAME}\"\n",
        "TOOLBOX_SERVICE_ACCOUNT_MEMBER = (\n",
        "    f\"serviceAccount:{TOOLBOX_SERVICE_ACCOUNT}@{PROJECT_ID}.iam.gserviceaccount.com\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DMLivT-MIcmV"
      },
      "source": [
        "## Getting Started With Gen AI Toolbox Locally\n",
        "\n",
        "Please follow the guide in https://github.com/googleapis/genai-toolbox/blob/main/docs/en/getting-started/colab_quickstart.ipynb\n",
        "\n",
        "This guide demonstrates how to quickly run [Toolbox](https://github.com/googleapis/genai-toolbox) end-to-end in Google Colab using Python, PostgreSQL, Google Gen AI and LangGraph or Llamaindex. Including:\n",
        "\n",
        "- Set up a `PostgreSQL database`.\n",
        "- Launch a Toolbox server.\n",
        "- Connect to Toolbox and develop a sample `Hotel Booking` agent.\n",
        "\n",
        "Here is the simplified flow of a Toolbox Application:\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/toolbox_flow.png\" alt=\"Toolbox Flow\"/>\n",
        "\n",
        "This notebook will reuse the example from the previous one, but demonstrate remote deployment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SZeYm8Vgmjjz"
      },
      "source": [
        "# Deploy PostgreSQL DB Remotely using Cloud SQL\n",
        "\n",
        "In the previous [toolbox quick-start notebook](https://github.com/googleapis/genai-toolbox/blob/main/docs/en/getting-started/colab_quickstart.ipynb), we demonstrated setting up a PostgreSQL database locally and populating it with data. In this notebook, we will show you how to deploy your PostgreSQL database remotely using Cloud SQL.\n",
        "\n",
        "If you previously learned how to create a PostgreSQL database using the guide at https://colab.sandbox.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_cloud_sql_pg_rag_agent.ipynb, feel free to skip `Step 1` and `Step 2` below, and proceed directly to `Step 3` to populate the data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QL58mPu9Hw7g"
      },
      "source": [
        "## Step 1: Before you begin\n",
        "\n",
        "1. In the Google Cloud console, on the project selector page, select or [create a Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects).\n",
        "1. [Make sure that billing is enabled for your Google Cloud project](https://cloud.google.com/billing/docs/how-to/verify-billing-enabled#console).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H3--jxtKiDR_"
      },
      "source": [
        "### Required roles\n",
        "\n",
        "To get the permissions that you need to complete the tutorial, ask your administrator to grant you the [Owner](https://cloud.google.com/iam/docs/understanding-roles#owner) (`roles/owner`) IAM role on your project. For more information about granting roles, see [Manage access](https://cloud.google.com/iam/docs/granting-changing-revoking-access)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9aGBuLA7aQ6O"
      },
      "source": [
        "### Define project information\n",
        "\n",
        "Initialize `gcloud` with your Project ID and resource location. At this time, only `us-central1` is supported."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vIeI4T_XVcDA"
      },
      "outputs": [],
      "source": [
        "!gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "veGoLZBYZjxY"
      },
      "source": [
        "### Enable APIs\n",
        "\n",
        "This tutorial uses the following billable components of Google Cloud, which you'll need to enable for this tutorial:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PcKjP3PiXDIi"
      },
      "outputs": [],
      "source": [
        "!gcloud services enable aiplatform.googleapis.com sqladmin.googleapis.com servicenetworking.googleapis.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S_yG0kddIvr7"
      },
      "source": [
        "## Step 2: Set up Cloud SQL\n",
        "\n",
        "Use the provided variable names or update the values to use a pre-existing Cloud SQL instance."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WwwSQ2ZUf51F"
      },
      "source": [
        "### Create a Cloud SQL instance\n",
        "\n",
        "This tutorial requires a Cloud SQL instance with public IP enabled."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LdbQ2Q0wdNiy"
      },
      "outputs": [],
      "source": [
        "# Create Cloud SQL instance\n",
        "!gcloud sql instances create {INSTANCE} \\\n",
        "  --database-version=POSTGRES_15 \\\n",
        "  --region={REGION} \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  --root-password={PASSWORD} \\\n",
        "  --cpu=1 \\\n",
        "  --memory=4GB \\\n",
        "  --assign-ip \\\n",
        "  --database-flags=cloudsql.iam_authentication=On"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZIjNIc71ixye"
      },
      "source": [
        "### Create a database\n",
        "\n",
        "Create a new database for the application using the Cloud SQL for LangChain library to establish a connection pool using the `PostgresEngine`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZnNy1NFVk9F7"
      },
      "source": [
        "By default, [IAM database authentication](https://cloud.google.com/sql/docs/mysql/iam-logins) will be used as the method of database authentication. This library uses the IAM principal belonging to the [Application Default Credentials (ADC)](https://cloud.google.com/docs/authentication/application-default-credentials) sourced from the environment.\n",
        "\n",
        "However, to smooth the onboarding process this tutorial will use the [built-in database authentication](https://cloud.google.com/sql/docs/mysql/built-in-authentication) using a username and password to access the Cloud SQL database can also be used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HiGfrIhY_wDq"
      },
      "outputs": [],
      "source": [
        "from sqlalchemy import text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pn4c8oZtixLI"
      },
      "outputs": [],
      "source": [
        "engine = await PostgresEngine.afrom_instance(\n",
        "    PROJECT_ID,\n",
        "    REGION,\n",
        "    INSTANCE,\n",
        "    database=DATABASE,\n",
        "    user=USER,\n",
        "    password=PASSWORD,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E8ih-faxCMMN"
      },
      "source": [
        "## Step 3: Populate Data\n",
        "\n",
        "In the previous [toolbox quick-start notebook](https://github.com/googleapis/genai-toolbox/blob/main/docs/en/getting-started/colab_quickstart.ipynb), we demonstrated how to populate data for the `Hotel Booking` Agent.\n",
        "\n",
        "We will reuse the sql query and populate data in Cloud SQL for PostgreSQL DB."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NOIp42SIAfoh"
      },
      "outputs": [],
      "source": [
        "async with engine._pool.connect() as conn:\n",
        "    await conn.execute(text(\"COMMIT\"))\n",
        "    await conn.execute(\n",
        "        text(\n",
        "            f\"\"\"\n",
        "            CREATE TABLE {TABLE_NAME}(\n",
        "              id            INTEGER NOT NULL PRIMARY KEY,\n",
        "              name          VARCHAR NOT NULL,\n",
        "              location      VARCHAR NOT NULL,\n",
        "              price_tier    VARCHAR NOT NULL,\n",
        "              checkin_date  DATE    NOT NULL,\n",
        "              checkout_date DATE    NOT NULL,\n",
        "              booked        BIT     NOT NULL\n",
        "            );\n",
        "            \"\"\"\n",
        "        )\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TkbX35pdBZJH"
      },
      "outputs": [],
      "source": [
        "async with engine._pool.connect() as conn:\n",
        "    await conn.execute(text(\"COMMIT\"))\n",
        "    await conn.execute(\n",
        "        text(\n",
        "            f\"\"\"\n",
        "            INSERT INTO {TABLE_NAME}(id, name, location, price_tier, checkin_date, checkout_date, booked)\n",
        "            VALUES\n",
        "              (1, 'Hilton Basel', 'Basel', 'Luxury', '2024-04-22', '2024-04-20', B'0'),\n",
        "              (2, 'Marriott Zurich', 'Zurich', 'Upscale', '2024-04-14', '2024-04-21', B'0'),\n",
        "              (3, 'Hyatt Regency Basel', 'Basel', 'Upper Upscale', '2024-04-02', '2024-04-20', B'0'),\n",
        "              (4, 'Radisson Lucerne', 'Lucerne', 'Midscale', '2024-04-24', '2024-04-05', B'0'),\n",
        "              (5, 'Best Western Bern', 'Bern', 'Upper Midscale', '2024-04-23', '2024-04-01', B'0'),\n",
        "              (6, 'InterContinental Geneva', 'Geneva', 'Luxury', '2024-04-23', '2024-04-28', B'0'),\n",
        "              (7, 'Sheraton Zurich', 'Zurich', 'Upper Upscale', '2024-04-27', '2024-04-02', B'0'),\n",
        "              (8, 'Holiday Inn Basel', 'Basel', 'Upper Midscale', '2024-04-24', '2024-04-09', B'0'),\n",
        "              (9, 'Courtyard Zurich', 'Zurich', 'Upscale', '2024-04-03', '2024-04-13', B'0'),\n",
        "              (10, 'Comfort Inn Bern', 'Bern', 'Midscale', '2024-04-04', '2024-04-16', B'0');\n",
        "            \"\"\"\n",
        "        )\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Dx2tzMLB1fn"
      },
      "outputs": [],
      "source": [
        "async with engine._pool.connect() as conn:\n",
        "    await conn.execute(text(\"COMMIT\"))\n",
        "    result = await conn.execute(text(f\"SELECT * from {TABLE_NAME};\"))\n",
        "    print(result.fetchall())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DEyqACccyGQR"
      },
      "source": [
        "Notice, the `Hotel Booking` agent will update the remote database, please run the following `repopulate_date()` before `repeated` run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VCSO89_hj9BU"
      },
      "outputs": [],
      "source": [
        "# @title `repopulate_data` function\n",
        "\n",
        "\n",
        "def repopulate_data():\n",
        "    async with engine._pool.connect() as conn:\n",
        "        await conn.execute(text(\"COMMIT\"))\n",
        "        await conn.execute(text(f\"DROP TABLE {TABLE_NAME};\"))\n",
        "\n",
        "    async with engine._pool.connect() as conn:\n",
        "        await conn.execute(text(\"COMMIT\"))\n",
        "        await conn.execute(\n",
        "            text(\n",
        "                f\"\"\"\n",
        "    CREATE TABLE {TABLE_NAME}(\n",
        "      id            INTEGER NOT NULL PRIMARY KEY,\n",
        "      name          VARCHAR NOT NULL,\n",
        "      location      VARCHAR NOT NULL,\n",
        "      price_tier    VARCHAR NOT NULL,\n",
        "      checkin_date  DATE    NOT NULL,\n",
        "      checkout_date DATE    NOT NULL,\n",
        "      booked        BIT     NOT NULL\n",
        "    );\n",
        "    \"\"\"\n",
        "            )\n",
        "        )\n",
        "\n",
        "    async with engine._pool.connect() as conn:\n",
        "        await conn.execute(text(\"COMMIT\"))\n",
        "        result = await conn.execute(text(f\"SELECT * from {TABLE_NAME};\"))\n",
        "        print(result.fetchall())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e9ZfPaG9FGj9"
      },
      "source": [
        "## Step 4: Create a user\n",
        "\n",
        "Set up the AI Platform Agent Engine Service Agent service account (`service-PROJECT_NUMBER@gcp-sa-aiplatform-re.iam.gserviceaccount.com`) as a database user - to log into the database, a database client - to connect to the database, and an AI Platform user - to connect to Vertex AI models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iH39YSf_iZle"
      },
      "outputs": [],
      "source": [
        "# Define service account\n",
        "PROJECT_NUMBER = !gcloud projects describe {PROJECT_ID} --format=\"value(projectNumber)\"\n",
        "SERVICE_ACCOUNT = f\"service-{PROJECT_NUMBER[0]}@gcp-sa-aiplatform-re.iam.gserviceaccount.com\"\n",
        "IAM_USER = SERVICE_ACCOUNT.replace(\".gserviceaccount.com\", \"\")\n",
        "\n",
        "# Force the creation of the AI Platform service accounts\n",
        "# The service accounts will be created at deploy time if not pre-created\n",
        "!gcloud beta services identity create --service=aiplatform.googleapis.com --project={PROJECT_ID}\n",
        "\n",
        "# Add a service account as database IAM user\n",
        "# For an IAM service account, supply the service account's address without the .gserviceaccount.com\n",
        "!gcloud sql users create {IAM_USER} \\\n",
        "  --instance={INSTANCE} \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  --type=cloud_iam_service_account\n",
        "\n",
        "# Grant IAM Permissions for database-user authentication\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "    --role=roles/cloudsql.instanceUser\n",
        "\n",
        "# Grant IAM permissions to access Cloud SQL instances\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "    --role=roles/cloudsql.client\n",
        "\n",
        "# Grant IAM permissions to access AI Platform services\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT}  \\\n",
        "    --role=roles/aiplatform.user\n",
        "\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT}  \\\n",
        "    --role=roles/serviceusage.serviceUsageConsumer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5GfAVv-E5i-f"
      },
      "outputs": [],
      "source": [
        "# Grant access of table to IAM users\n",
        "engine = await PostgresEngine.afrom_instance(\n",
        "    PROJECT_ID, REGION, INSTANCE, DATABASE, user=USER, password=PASSWORD\n",
        ")\n",
        "\n",
        "async with engine._pool.connect() as conn:\n",
        "    await conn.execute(text(f'GRANT SELECT ON {TABLE_NAME} TO \"{IAM_USER}\";'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PprJNEVhmUvQ"
      },
      "source": [
        "# Deploy Toolbox to Cloud Run Remotely"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mRBlMDHrwlJ1"
      },
      "source": [
        "Please follow the guide in https://googleapis.github.io/genai-toolbox/how-to/deploy_toolbox/.\n",
        "\n",
        "This guide demonstrates how to deploy your toolbox to Cloud Run."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x1Kph-Ybx053"
      },
      "source": [
        "## Step 1: Enable APIs\n",
        "This tutorial uses the following billable components of Google Cloud, which you'll need to enable for this tutorial:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lpb_0Y4EKacZ"
      },
      "outputs": [],
      "source": [
        "!gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_9abn8pVvk4G"
      },
      "outputs": [],
      "source": [
        "!gcloud services enable run.googleapis.com \\\n",
        "  cloudbuild.googleapis.com \\\n",
        "  artifactregistry.googleapis.com \\\n",
        "  iam.googleapis.com \\\n",
        "  secretmanager.googleapis.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0oT_q_Ytv3aD"
      },
      "source": [
        "To create an IAM account, you must have the following IAM permissions (or roles):\n",
        "\n",
        "    Create Service Account role (roles/iam.serviceAccountCreator)\n",
        "\n",
        "To create a secret, you must have the following roles:\n",
        "\n",
        "    Secret Manager Admin role (roles/secretmanager.admin)\n",
        "\n",
        "To deploy to Cloud Run, you must have the following set of roles:\n",
        "\n",
        "    Cloud Run Developer (roles/run.developer)\n",
        "    Service Account User role (roles/iam.serviceAccountUser)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xc4XGVy_vuBt"
      },
      "outputs": [],
      "source": [
        "!gcloud iam service-accounts create {TOOLBOX_SERVICE_ACCOUNT}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jODb3pT1wpf0"
      },
      "outputs": [],
      "source": [
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member {TOOLBOX_SERVICE_ACCOUNT_MEMBER} \\\n",
        "    --role roles/secretmanager.secretAccessor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WeNktwuDCCNy"
      },
      "outputs": [],
      "source": [
        "!gcloud beta run services add-iam-policy-binding \\\n",
        "    --region=us-central1 \\\n",
        "    --member=allUsers \\\n",
        "    --role=roles/run.invoker toolbox"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YNvxRtxS-EOc"
      },
      "outputs": [],
      "source": [
        "# Grant IAM Permissions for database-user authentication\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member={TOOLBOX_SERVICE_ACCOUNT_MEMBER} \\\n",
        "    --role=roles/cloudsql.instanceUser\n",
        "\n",
        "# Grant the Toolbox Service Account permission to connect to the Cloud SQL instance\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member={TOOLBOX_SERVICE_ACCOUNT_MEMBER} \\\n",
        "    --role=roles/cloudsql.client"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NG8320QSyoNv"
      },
      "source": [
        "## Step 2: Create the configure for Toolbox\n",
        "\n",
        "Following the schema defined at https://googleapis.github.io/genai-toolbox/resources/sources/, create a Toolbox configuration file (e.g., tools.yml)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "72B7AxvmCeC5"
      },
      "outputs": [],
      "source": [
        "# Create a tools file at runtime.\n",
        "# You can also upload a tools file and use that to run toolbox.\n",
        "tools_file_name = \"tools.yml\"\n",
        "file_content = rf\"\"\"\n",
        "sources:\n",
        "  my-cloud-sql-source:\n",
        "    kind: cloud-sql-postgres\n",
        "    project: {PROJECT_ID}\n",
        "    region: {REGION}\n",
        "    instance: {INSTANCE}\n",
        "    database: {DATABASE}\n",
        "    user: {USER}\n",
        "    password: {PASSWORD}\n",
        "\n",
        "\n",
        "tools:\n",
        "  search-hotels-by-name:\n",
        "    kind: postgres-sql\n",
        "    source: my-cloud-sql-source\n",
        "    description: Search for hotels based on name.\n",
        "    parameters:\n",
        "      - name: name\n",
        "        type: string\n",
        "        description: The name of the hotel.\n",
        "    statement: SELECT * FROM hotels WHERE name ILIKE '%' || \\$1 || '%';\n",
        "  search-hotels-by-location:\n",
        "    kind: postgres-sql\n",
        "    source: my-cloud-sql-source\n",
        "    description: Search for hotels based on location.\n",
        "    parameters:\n",
        "      - name: location\n",
        "        type: string\n",
        "        description: The location of the hotel.\n",
        "    statement: SELECT * FROM hotels WHERE location ILIKE '%' || \\$1 || '%';\n",
        "  book-hotel:\n",
        "    kind: postgres-sql\n",
        "    source: my-cloud-sql-source\n",
        "    description: >-\n",
        "       Book a hotel by its ID. If the hotel is successfully booked, returns a NULL, raises an error if not.\n",
        "    parameters:\n",
        "      - name: hotel_id\n",
        "        type: string\n",
        "        description: The ID of the hotel to book.\n",
        "    statement: UPDATE hotels SET booked = B'1' WHERE id = \\$1;\n",
        "  update-hotel:\n",
        "    kind: postgres-sql\n",
        "    source: my-cloud-sql-source\n",
        "    description: >-\n",
        "      Update a hotel's check-in and check-out dates by its ID. Returns a message\n",
        "      indicating  whether the hotel was successfully updated or not.\n",
        "    parameters:\n",
        "      - name: hotel_id\n",
        "        type: string\n",
        "        description: The ID of the hotel to update.\n",
        "      - name: checkin_date\n",
        "        type: string\n",
        "        description: The new check-in date of the hotel.\n",
        "      - name: checkout_date\n",
        "        type: string\n",
        "        description: The new check-out date of the hotel.\n",
        "    statement: >-\n",
        "      UPDATE hotels SET checkin_date = CAST(\\$2 as date), checkout_date = CAST(\\$3\n",
        "      as date) WHERE id = \\$1;\n",
        "  cancel-hotel:\n",
        "    kind: postgres-sql\n",
        "    source: my-cloud-sql-source\n",
        "    description: Cancel a hotel by its ID.\n",
        "    parameters:\n",
        "      - name: hotel_id\n",
        "        type: string\n",
        "        description: The ID of the hotel to cancel.\n",
        "    statement: UPDATE hotels SET booked = B'0' WHERE id = \\$1;\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IXsRLYDDFJSV"
      },
      "outputs": [],
      "source": [
        "# Write the file content into the tools file.\n",
        "! echo \"{file_content}\" > \"{tools_file_name}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ANQgzG0oFJSW"
      },
      "outputs": [],
      "source": [
        "TOOLS_FILE_PATH = f\"/content/{tools_file_name}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MWgVrJpxzza_"
      },
      "source": [
        "## Step 3: Deploy to Cloud Run\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "puRKGqaUz1id"
      },
      "source": [
        "Upload `tools.yaml` as a secret:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z1f6NMN8z1MQ"
      },
      "outputs": [],
      "source": [
        "!gcloud secrets create tools --data-file={TOOLS_FILE_PATH}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_aBzfUuW0gdo"
      },
      "source": [
        "Set an environment variable to the container image that you want to use for cloud run:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_ANyWeNrwsYQ"
      },
      "outputs": [],
      "source": [
        "IMAGE = \"us-central1-docker.pkg.dev/database-toolbox/toolbox/toolbox:latest\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NSeKCQ_N0nV5"
      },
      "source": [
        "Deploy Toolbox to Cloud Run using the following command:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BXn9Y6lM0kvp"
      },
      "outputs": [],
      "source": [
        "!gcloud run deploy toolbox \\\n",
        "    --image {IMAGE} \\\n",
        "    --service-account toolbox-identity \\\n",
        "    --region us-central1 \\\n",
        "    --set-secrets \"/app/tools.yaml=tools:latest\" \\\n",
        "    --args=\"--tools_file=/app/tools.yaml\",\"--address=0.0.0.0\",\"--port=8080\" \\\n",
        "    --allow-unauthenticated"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2z1zQyjn-mds"
      },
      "source": [
        "## Step 4: Connecting with Toolbox Client SDK\n",
        "\n",
        "You should get a service URL like:\n",
        "https://toolbox-[project-number].us-central1.run.app\n",
        "\n",
        "If not, run the following to get the url:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kgtV334e-sHe"
      },
      "outputs": [],
      "source": [
        "!gcloud run services describe toolbox --format 'value(status.url)'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Eh1FesN5IdjV"
      },
      "outputs": [],
      "source": [
        "# @title Set up Toolbox Service URL Parameters\n",
        "# fmt: off\n",
        "TOOLBOX_ENDPOINT = \"[your-toolbox-service-url]\"  # @param {type:\"string\"} Your Toolbox Service URL.\n",
        "# fmt: on"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "148gfbHlJDOC"
      },
      "source": [
        "Test the connection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XV6tunIoAju_"
      },
      "outputs": [],
      "source": [
        "from toolbox_langchain import ToolboxClient\n",
        "\n",
        "toolbox = ToolboxClient(TOOLBOX_ENDPOINT)\n",
        "toolbox.load_toolset()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pgbn_k1nBm_P"
      },
      "source": [
        "# Develop and Test your Agent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uepz1IGYMdIS"
      },
      "source": [
        "## Step 1: Develop the Hotel Booking Agent\n",
        "\n",
        "\n",
        "This section follows the [MCP toolbox for databases quick-start](https://github.com/googleapis/genai-toolbox/blob/main/docs/en/getting-started/colab_quickstart.ipynb) to build a `hotel booking agent`, and demonstrates how to wrap your hotel booking agent into a [Python class that is deployable on Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LWo0K0K-MQ31"
      },
      "source": [
        "First, we define the class as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kg88Iwd6EXNg"
      },
      "outputs": [],
      "source": [
        "class HotelBookingAgent:\n",
        "    def __init__(\n",
        "        self, model: str, project: str, location: str, toolbox_endpoint: str\n",
        "    ) -> None:\n",
        "        self.model = model\n",
        "        self.project_id = project\n",
        "        self.location = location\n",
        "        self.toolbox_endpoint = toolbox_endpoint\n",
        "        self.prompt = \"\"\"\n",
        "          You're a helpful hotel assistant. You handle hotel searching, booking and\n",
        "          cancellations. When the user searches for a hotel, mention its name, id,\n",
        "          location and price tier. Always mention hotel id while performing any\n",
        "          searches. This is very important for any operations. For any bookings or\n",
        "          cancellations, please provide the appropriate confirmation. Be sure to\n",
        "          update checkin or checkout dates if mentioned by the user.\n",
        "          Don't ask for confirmations from the user.\n",
        "        \"\"\"\n",
        "\n",
        "    # The set_up method is used to define application initialization logic\n",
        "    def set_up(self) -> None:\n",
        "        from langchain_google_vertexai import ChatVertexAI\n",
        "        from langgraph.checkpoint.memory import MemorySaver\n",
        "        from langgraph.prebuilt import create_react_agent\n",
        "        from toolbox_langchain import ToolboxClient\n",
        "\n",
        "        # Create an LLM to bind with the agent.\n",
        "        model = ChatVertexAI(model_name=self.model, project=self.project_id)\n",
        "\n",
        "        # Load the tools from the Toolbox server\n",
        "        client = ToolboxClient(self.toolbox_endpoint)\n",
        "        tools = client.load_toolset()\n",
        "\n",
        "        # Create a ReAct agent\n",
        "        self.runnable = create_react_agent(model, tools, checkpointer=MemorySaver())\n",
        "\n",
        "    # The query method will be used to send inputs to the agent\n",
        "    def query(self, message: str):\n",
        "        \"\"\"Query the application.\n",
        "\n",
        "        Args:\n",
        "            message: The user message.\n",
        "\n",
        "        Returns:\n",
        "            str: The LLM response.\n",
        "        \"\"\"\n",
        "        config = {\"configurable\": {\"thread_id\": \"thread-1\"}}\n",
        "        inputs = {\"messages\": [(\"user\", self.prompt + message)]}\n",
        "        response = self.runnable.invoke(inputs, stream_mode=\"values\", config=config)\n",
        "\n",
        "        return response[\"messages\"][-1].content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bCOszwwTMTou"
      },
      "source": [
        "Next, we initialize the agent as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tLh4lr-gcjLJ"
      },
      "outputs": [],
      "source": [
        "agent = HotelBookingAgent(\n",
        "    model=\"gemini-2.0-flash\",\n",
        "    toolbox_endpoint=TOOLBOX_ENDPOINT,\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xOvFQxcsMYoi"
      },
      "source": [
        "Finally, we set up the agent before testing it locally:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r2XQDOW1cjLJ"
      },
      "outputs": [],
      "source": [
        "agent.set_up()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j4lIXO7stTuA"
      },
      "source": [
        "## Step 2: Test the agent locally"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sWvPnkQ5mggj"
      },
      "source": [
        "Your local `Hotel Booking` Agent will process user's query, make tool calls to the remote `Toolbox` service, which in turn will execute SQL queries against the remote `Cloud SQL` database."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Sk6T_DacjLJ"
      },
      "outputs": [],
      "source": [
        "queries = [\n",
        "    \"Find hotels in Basel with Basel in its name.\",\n",
        "    \"Can you book the Hilton Basel for me?\",\n",
        "    \"Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead.\",\n",
        "    \"My check in dates would be from April 10, 2024 to April 19, 2024.\",\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GhBkFndebwxn"
      },
      "source": [
        "Notice, the following command will update the remote database, if you want to repeat your query, please run `repopulate_date()` before repeated run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HDB9DUD8cjLJ"
      },
      "outputs": [],
      "source": [
        "for query in queries:\n",
        "    response = agent.query(message=query)\n",
        "    print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6f2cc9278b8e"
      },
      "source": [
        "Expected output:\n",
        "```\n",
        "I found the following hotels in Basel: Hilton Basel (ID: 1), Hyatt Regency Basel (ID: 3), and Holiday Inn Basel (ID: 8).\n",
        "\n",
        "OK. I have booked the Hilton Basel (ID: 1) for you.\n",
        "\n",
        "OK. I have cancelled the Hilton Basel (ID: 1) and booked the Hyatt Regency (ID: 3) for you.\n",
        "\n",
        "OK. I have updated your check-in date to April 10, 2024 and check-out date to April 19, 2024 for the Hyatt Regency (ID: 3).\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AZYIlDvlSeoA"
      },
      "source": [
        "# Deploy your agent to Agent Engine Remotely\n",
        "\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",
        "Here, you'll use the Vertex AI SDK for Agent Engine, which brings together the model, tools, and reasoning that you've built up so far."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TqSfLzpdNG9J"
      },
      "source": [
        "## Step 1: Create a Cloud Storage bucket\n",
        "\n",
        "Create or reuse and existing Cloud Storage bucket. Agent Engine stages the artifacts of your applications in a Cloud Storage bucket as part of the deployment process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sptkevO4aUT1"
      },
      "outputs": [],
      "source": [
        "# Create a Cloud Storage bucket, if it doesn't already exist\n",
        "!gsutil mb -c standard {STAGING_BUCKET}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TCIC2oVTVkLh"
      },
      "source": [
        "## Step 2: Initialize the Vertex AI SDK for Agent Engine\n",
        "Initialize the Vertex AI SDK for Agent Engine with your project, location, and Cloud Storage bucket.\n",
        "Check out the https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/set-up#sdk for more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M5sF00H3VjoG"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "from google.colab import auth\n",
        "from vertexai import agent_engines\n",
        "\n",
        "auth.authenticate_user()\n",
        "\n",
        "vertexai.init(\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I_hlYrmcVWqp"
      },
      "source": [
        "## Step 3: Deploy your agent remotely\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!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k2nGSr2_JWcc"
      },
      "outputs": [],
      "source": [
        "remote_agent = agent_engines.create(\n",
        "    HotelBookingAgent(\n",
        "        model=\"gemini-2.0-flash\",\n",
        "        toolbox_endpoint=TOOLBOX_ENDPOINT,\n",
        "        project=PROJECT_ID,\n",
        "        location=REGION,\n",
        "    ),\n",
        "    requirements=[\n",
        "        \"toolbox-langchain==0.1.0\",\n",
        "        \"google-cloud-aiplatform[agent_engines,langchain]==1.87.0\",\n",
        "    ],\n",
        "    display_name=\"HotelBookingAgent\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ebjYuN6mFRC"
      },
      "source": [
        "## Step 4: Test the Deployed Agent.\n",
        "\n",
        "Now you have all `PostgreSQL DB`, `Toolbox` and `Hotel Booking Agent` deployed remotely!\n",
        "\n",
        "Your remote `Hotel Booking` Agent will process user's query, make tool calls to the remote `Toolbox` service, which in turn will execute SQL queries against the remote `Cloud SQL` database."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uZlHcNZedmNs"
      },
      "source": [
        "Notice, the following command will update the remote database, please call `repopulate_date()` before repeated run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aIBC6CgRdJ3C"
      },
      "outputs": [],
      "source": [
        "for query in queries:\n",
        "    response = remote_agent.query(message=query)\n",
        "    print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ce30c1cb84f4"
      },
      "source": [
        "Expected output:\n",
        "\n",
        "```\n",
        "I found the following hotels in Basel: Hilton Basel (ID: 1), Hyatt Regency Basel (ID: 3), and Holiday Inn Basel (ID: 8).\n",
        "\n",
        "OK. I have booked the Hilton Basel (ID: 1) for you.\n",
        "\n",
        "OK. I have cancelled the Hilton Basel (ID: 1) and booked the Hyatt Regency (ID: 3) for you.\n",
        "\n",
        "OK. I have updated your check-in date to April 10, 2024 and check-out date to April 19, 2024 for the Hyatt Regency (ID: 3).\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9f2f7d3ed7bd"
      },
      "source": [
        "### Querying your deployed agent\n",
        "\n",
        "You've now deployed your 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.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": null,
      "metadata": {
        "id": "78af4442827e"
      },
      "outputs": [],
      "source": [
        "# from vertexai 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 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": "6415b2a47a0e"
      },
      "source": [
        "# Cleaning up\n",
        "\n",
        "After you've finished, 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": "markdown",
      "metadata": {
        "id": "tBc48ZHOJS6J"
      },
      "source": [
        "### Deleting the project\n",
        "\n",
        "The easiest way to eliminate billing is to delete the project that you created for the tutorial.\n",
        "\n",
        "1. In the Google Cloud console, go to the [Manage resources](https://console.cloud.google.com/iam-admin/projects?_ga=2.235586881.1783688455.1719351858-1945987529.1719351858) page.\n",
        "1. In the project list, select the project that you want to delete, and then click Delete.\n",
        "1. In the dialog, type the project ID, and then click Shut down to delete the project.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ed-BFtW-JPbI"
      },
      "source": [
        "### Deleting tutorial resources\n",
        "\n",
        "Delete the Agent Engine instance(s) and Cloud SQL instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LgNlHrxkb6c-"
      },
      "outputs": [],
      "source": [
        "# Delete the Agent Engine instance\n",
        "remote_agent.delete()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "goyrqS2_I8Hs"
      },
      "outputs": [],
      "source": [
        "# Or delete all Agent Engine apps\n",
        "apps = remote_agent.list()\n",
        "for app in apps:\n",
        "    app.delete()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "afxupsbwpaA8"
      },
      "outputs": [],
      "source": [
        "# Delete the Toolbox Cloud Run instance\n",
        "!gcloud run services delete toolbox --region={REGION} --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "odvj8aKpb3Wi"
      },
      "outputs": [],
      "source": [
        "# Delete the Cloud SQL instance\n",
        "!gcloud sql instances delete {INSTANCE} \\\n",
        "  --project={PROJECT_ID}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tutorial_mcp_toolbox_for_databases.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
