{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Vs2C5ZBkrjmrHTfAyz77hpV"
      },
      "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": "VoJZgH6X77se"
      },
      "source": [
        "# MCP Server with Gemini Enterprise"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qfDq3m0R8Fe1"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/search/gemini-enterprise/mcp_with_gemini_enterprise.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%2Fsearch%2Fgemini-enterprise%2Fmcp_with_gemini_enterprise.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/search/gemini-enterprise/mcp_with_gemini_enterprise.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/search/gemini-enterprise/mcp_with_gemini_enterprise.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "\n",
        "\n",
        "<br>\n",
        "<br>\n",
        "<br>\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/search/gemini-enterprise/mcp_with_gemini_enterprise.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/search/gemini-enterprise/mcp_with_gemini_enterprise.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/search/gemini-enterprise/mcp_with_gemini_enterprise.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/search/gemini-enterprise/mcp_with_gemini_enterprise.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/search/gemini-enterprise/mcp_with_gemini_enterprise.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": "UD1Cy9O88Zn9"
      },
      "source": [
        "| Authors |\n",
        "| --- |\n",
        "| Parag Mhatre |\n",
        "| Sharmila Devi |\n",
        "| Rupjit Chakraborty|"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "28b524cdd168"
      },
      "source": [
        "## Overview : MCP Server with Gemini Enterprise\n",
        "\n",
        "This notebook demonstrates a complete end-to-end integration of a business application with Gemini Enterprise, leveraging a Model Context Protocol (MCP) server built with FastMCP and an agent developed using the Agent Development Kit (ADK). This guide is designed to assist developers in understanding and implementing similar solutions for their upcoming projects.\n",
        "\n",
        "### Background\n",
        "The increasing demand for integrating business applications with intelligent agents highlights the need for robust and scalable solutions. This notebook addresses this by providing a practical demonstration of how an MCP server can be seamlessly integrated with Gemini Enterprise using the ADK.\n",
        "\n",
        "### Business Application Scenario\n",
        "For this demonstration, we use a common business application: employee leave application. The application requires employees to provide their employee ID, start date, and end date for leave requests.\n",
        "\n",
        "### Notebook Overview\n",
        "This notebook guides you through the following key steps:\n",
        "\n",
        "#### 1. BigQuery Setup\n",
        "\n",
        "> The process begins with setting up the necessary data infrastructure. You'll learn how to: <br>\n",
        "Create a BigQuery dataset. <br>\n",
        "Create a BigQuery table to store employee leave application data.\n",
        "\n",
        "#### 2. MCP Server Deployment to Cloud Run\n",
        "> This section focuses on deploying the MCP server, which acts as the bridge between your business application and the agents. <br>\n",
        "The current MCP server is built using the FastMCP framework, providing a robust RESTful API for communication. <br>\n",
        "You'll be guided through the process of building a container image for the MCP server. <br>\n",
        "Instructions for deploying this container to Cloud Run will be provided, ensuring a scalable and serverless deployment.\n",
        "\n",
        "#### 3. ADK Configuration and MCP Server Communication\n",
        "> Once the MCP server is deployed, the notebook demonstrates how to configure the ADK to interact with it. <br>\n",
        "Learn how to set up your ADK environment. <br>\n",
        "Understand how to use MCP Tools to facilitate communication between your ADK-based agent and the deployed MCP server. <br>\n",
        "This section includes steps for testing the integration to ensure smooth data exchange.\n",
        "\n",
        "#### 4. Agent Deployment to Agent Engine\n",
        "> After successful testing, the notebook walks you through deploying your ADK-built agent. <br>\n",
        "Detailed steps are provided for deploying the agent to the Agent Engine.\n",
        "\n",
        "#### 5. Agent Engine Registration with Gemini Enterprise\n",
        "> The final step involves integrating the deployed agent with the broader Gemini Enterprise ecosystem.\n",
        "\n",
        "Learn how to register the Agent Engine with Gemini Enterprise using REST APIs, making your agent discoverable and accessible within the Gemini Enterprise environment.\n",
        "\n",
        "\n",
        "**Open Issue:**\n",
        "> 1. ADK currently utilizes the MCP Tool to interface with the MCP server, which is deployed on Cloud Run. Authorization headers for Cloud Run expire after one hour. Efforts are underway to integrate the Agent Engine service account for interaction with Cloud Run."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k0nKem_m_tOt"
      },
      "source": [
        "### MCP Server for Human Resources Management Application"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5bb4021007d1"
      },
      "outputs": [],
      "source": [
        "# TODO for Developer: Update project id.\n",
        "PROJECT_NUMBER = \"[your-project-number]\"\n",
        "\n",
        "# TODO for Developer: Update project name.\n",
        "PROJECT_ID = \"[your-project-id]\"\n",
        "\n",
        "DATASET_ID = \"bqds_hr_data\"\n",
        "TABLE_NAME = \"employee_leave\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bc0a45394998"
      },
      "outputs": [],
      "source": [
        "# Setup Google Cloud Product project.\n",
        "!gcloud config set project {PROJECT_ID}\n",
        "!gcloud config get-value project\n",
        "\n",
        "# Enable required services.\n",
        "!gcloud services enable iam.googleapis.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K3QuYpHs_wR0"
      },
      "source": [
        "#### Step 1: BigQuery Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a587b1e6b850"
      },
      "outputs": [],
      "source": [
        "# Create required BigQuery dataset.\n",
        "!bq mk --dataset --PROJECT_NUMBER {PROJECT_NUMBER} {DATASET_ID}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a4f38854a72f"
      },
      "outputs": [],
      "source": [
        "# Create required BigQuery table.\n",
        "!bq mk \\\n",
        "  --table \\\n",
        "  {PROJECT_NUMBER}:{DATASET_ID}.{TABLE_NAME} \\\n",
        "  employee_id:INTEGER,leave_start_date:DATE,leave_end_date:DATE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "af2224878401"
      },
      "outputs": [],
      "source": [
        "# Setting up environment variable.\n",
        "%env PROJECT_ID={PROJECT_ID}\n",
        "%env DATASET_ID={DATASET_ID}\n",
        "%env TABLE_NAME={TABLE_NAME}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0d5680e059dd"
      },
      "source": [
        "#### Step 2: MCP Server Deployment to Cloud Run"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "79ae8974bb0b"
      },
      "source": [
        "##### Step 2.1 : Build MCP Server container image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9b93783e1936"
      },
      "outputs": [],
      "source": [
        "%%writefile mcp_server.py\n",
        "import os\n",
        "import json\n",
        "import uvicorn\n",
        "\n",
        "from requests.exceptions import RequestException\n",
        "from starlette.applications import Starlette\n",
        "from starlette.requests import Request\n",
        "from starlette.routing import Route, Mount\n",
        "\n",
        "from mcp.server.fastmcp import FastMCP\n",
        "from mcp.shared.exceptions import McpError\n",
        "from mcp.types import ErrorData, INTERNAL_ERROR, INVALID_PARAMS\n",
        "from mcp.server.sse import SseServerTransport\n",
        "from google.cloud import bigquery\n",
        "\n",
        "# Create an MCP server instance with an identifier (\"HR Tool\")\n",
        "mcp = FastMCP(\"HR Tool\")\n",
        "APP_HOST = os.environ.get(\"APP_HOST\", \"0.0.0.0\")\n",
        "APP_PORT = os.environ.get(\"APP_PORT\", 8080)\n",
        "\n",
        "# Retrieve environment variables.\n",
        "PROJECT_NUMBER = os.environ.get(\"PROJECT_ID\")\n",
        "dataset_id = os.environ.get(\"DATASET_ID\")\n",
        "table_name = os.environ.get(\"TABLE_NAME\")\n",
        "table_id = f\"{PROJECT_NUMBER}.{dataset_id}.{table_name}\"\n",
        "\n",
        "# Setup BigQuery client.\n",
        "client = bigquery.Client(project=PROJECT_NUMBER)\n",
        "\n",
        "\n",
        "@mcp.tool()\n",
        "def apply_leave(\n",
        "    employee_id: int, start_date: str, end_date: str\n",
        ") -> bigquery.job.LoadJob:\n",
        "    \"\"\"\n",
        "    Inserts data into a BigQuery table.\n",
        "\n",
        "    Args:\n",
        "        employee_id: Employee identification int. should be in string format.\n",
        "        start_date: Start date of a leave. Dates should be in 'YYYY-MM-DD' string format.\n",
        "        end_date: End date of a leave. Dates should be in 'YYYY-MM-DD' string format.\n",
        "\n",
        "    Returns:\n",
        "        A BigQuery LoadJob object if the insert operation is successful.\n",
        "    \"\"\"\n",
        "\n",
        "    # Prepare row to be inserted.\n",
        "    data_to_insert = [\n",
        "        {\n",
        "            \"employee_id\": employee_id,\n",
        "            \"leave_start_date\": start_date,\n",
        "            \"leave_end_date\": end_date,\n",
        "        },\n",
        "    ]\n",
        "\n",
        "    # insert row to BigQuery.\n",
        "    errors = client.insert_rows_json(table_id, data_to_insert)\n",
        "\n",
        "    if errors:\n",
        "        print(f\"Encountered errors while inserting rows: {errors}\")\n",
        "        raise Exception(f\"Failed to insert rows: {errors}\")\n",
        "    else:\n",
        "        print(f\"Successfully inserted {len(data_to_insert)} rows into {table_id}\")\n",
        "        return \"Leave applied successfuly. for \" + json.dumps(data_to_insert)\n",
        "\n",
        "\n",
        "sse = SseServerTransport(\"/messages/\")\n",
        "\n",
        "\n",
        "async def handle_sse(request: Request) -> None:\n",
        "    _server = mcp._mcp_server\n",
        "    async with sse.connect_sse(\n",
        "        request.scope,\n",
        "        request.receive,\n",
        "        request._send,\n",
        "    ) as (reader, writer):\n",
        "        await _server.run(reader, writer, _server.create_initialization_options())\n",
        "\n",
        "\n",
        "app = Starlette(\n",
        "    debug=True,\n",
        "    routes=[\n",
        "        Route(\"/sse\", endpoint=handle_sse),\n",
        "        Mount(\"/messages/\", app=sse.handle_post_message),\n",
        "    ],\n",
        ")\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    uvicorn.run(app, host=APP_HOST, port=APP_PORT)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "68ed7244d84a"
      },
      "outputs": [],
      "source": [
        "%%writefile Dockerfile\n",
        "# Use an official Python runtime as a parent image\n",
        "FROM python:3.10-slim\n",
        "\n",
        "# Set the working directory in the container\n",
        "WORKDIR /app\n",
        "\n",
        "# --- Dependency Installation ---\n",
        "COPY requirements.txt /app/requirements.txt\n",
        "RUN pip install --no-cache-dir -r requirements.txt\n",
        "\n",
        "# --- Application Code ---\n",
        "COPY . /app\n",
        "\n",
        "# --- Environment ---\n",
        "ENV PYTHONPATH = /app\n",
        "\n",
        "# Make port 8080 available to the world outside this container\n",
        "# Cloud Run uses the PORT env var, but EXPOSE is good practice.\n",
        "EXPOSE 8080\n",
        "\n",
        "# --- Run the application ---\n",
        "CMD [\"python\", \"mcp_server.py\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2c3d35247c0b"
      },
      "outputs": [],
      "source": [
        "%%writefile requirements.txt\n",
        "google-adk == 0.3.0\n",
        "mcp[cli] == 1.5.0\n",
        "requests == 2.32.3\n",
        "google-cloud-bigquery ==3.35.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0a0ec35d06e4"
      },
      "source": [
        "##### Step 2.1: Deploy MCP Server to Cloud Run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d370f0eb341e"
      },
      "outputs": [],
      "source": [
        "# Setup Image repository URL\n",
        "IMAGE_TAG = \"latest\"\n",
        "MCP_IMAGE_NAME = \"hr-tool-server\"\n",
        "REPO_NAME = \"mcp-server\"\n",
        "REGION = \"us-central1\"\n",
        "IMAGE_PATH = (\n",
        "    f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPO_NAME}/{MCP_IMAGE_NAME}:{IMAGE_TAG}\"\n",
        ")\n",
        "SERVICE_NAME = \"hr-tool-server\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bbc9198085e4"
      },
      "outputs": [],
      "source": [
        "# Create repository\n",
        "!gcloud artifacts repositories create {REPO_NAME} \\\n",
        "  --repository-format=docker \\\n",
        "  --location={REGION} \\\n",
        "  --description=\"Repository for my application's Docker images\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "153e3ac7f8d0"
      },
      "outputs": [],
      "source": [
        "# Build image and push to artifact registry.\n",
        "!gcloud builds submit . \\\n",
        "  --tag={IMAGE_PATH} \\\n",
        "  --project={PROJECT_ID}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cba3beca8369"
      },
      "outputs": [],
      "source": [
        "# Deploy MCP Server to cloud run.\n",
        "cmd = f\"\"\"gcloud run deploy {SERVICE_NAME} \\\n",
        "  --image={IMAGE_PATH} \\\n",
        "  --platform=managed \\\n",
        "  --region={REGION} \\\n",
        "  --no-allow-unauthenticated \\\n",
        "  --set-env-vars=\"APP_HOST=0.0.0.0\" \\\n",
        "  --set-env-vars=\"APP_PORT=8080\" \\\n",
        "  --set-env-vars=\"GOOGLE_GENAI_USE_VERTEXAI=TRUE\" \\\n",
        "  --set-env-vars=\"GOOGLE_CLOUD_LOCATION={REGION}\" \\\n",
        "  --set-env-vars=\"GOOGLE_CLOUD_PROJECT={PROJECT_NUMBER}\" \\\n",
        "  --set-env-vars=\"PROJECT_ID={PROJECT_ID}\" \\\n",
        "  --set-env-vars=\"DATASET_ID={DATASET_ID}\" \\\n",
        "  --set-env-vars=\"TABLE_NAME={TABLE_NAME}\" \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  --min-instances=1\"\"\"\n",
        "!{cmd}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2e37d2d52406"
      },
      "source": [
        "#### Step 3: ADK Configuration and MCP Server Communication"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dea5bf7484e1"
      },
      "outputs": [],
      "source": [
        "pip install --quiet google-adk==1.4.2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "97e86592081a"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "\n",
        "import google.auth.transport.requests\n",
        "import google.oauth2.id_token\n",
        "from dotenv import load_dotenv\n",
        "from google.adk.agents import Agent\n",
        "from google.adk.tools.mcp_tool.mcp_session_manager import SseServerParams\n",
        "from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset\n",
        "from google.genai import types\n",
        "from vertexai.preview import reasoning_engines"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "1dbb9681ee2b"
      },
      "outputs": [],
      "source": [
        "# The URL of the target Cloud Run service\n",
        "MCP_SERVER_URL = \"https://hr-tool-server-[your-project-number].us-central1.run.app/sse\"  # TODO for Developer : Copy Paste the url you will receive after cloud run deployment.\n",
        "\n",
        "# Create a Google Auth request object\n",
        "auth_req = google.auth.transport.requests.Request()\n",
        "\n",
        "# Generate the ID token\n",
        "id_token = google.oauth2.id_token.fetch_id_token(auth_req, MCP_SERVER_URL)\n",
        "\n",
        "# Make the request with the ID token\n",
        "headers = {\"Authorization\": f\"Bearer {id_token}\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b69ed8e2715b"
      },
      "source": [
        "##### Step 3.1: Register MCP Tool to ADK Agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "bc754ccfda73"
      },
      "outputs": [],
      "source": [
        "agent_prompt = \"\"\"\n",
        "You are having tools to apply leave. If your user is looking for applying leave then ask for employee id, start date and end date. once you have all the information then apply leave.\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "c5b7fcd20d13"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "INFO:__main__:Initializing agent...\n",
            "INFO:__main__:MCP Toolset created successfully.\n",
            "INFO:__main__:Agent initialized successfully.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Attempting to connect to MCP Filesystem server...\n",
            "LlmAgent created.\n"
          ]
        }
      ],
      "source": [
        "# Load environment variables from .env file in the parent directory\n",
        "# Place this near the top, before using env vars like API keys\n",
        "load_dotenv()\n",
        "\n",
        "logging.basicConfig(level=logging.INFO)\n",
        "log = logging.getLogger(__name__)\n",
        "\n",
        "# --- Global variables ---\n",
        "# Define them first, initialize as None\n",
        "root_agent: Agent | None = None\n",
        "\n",
        "\n",
        "def get_tools_async():\n",
        "    print(\"Attempting to connect to MCP Filesystem server...\")\n",
        "    tools = MCPToolset(\n",
        "        connection_params=SseServerParams(url=MCP_SERVER_URL, headers=headers),\n",
        "        errlog=None,\n",
        "    )\n",
        "    log.info(\"MCP Toolset created successfully.\")\n",
        "\n",
        "    return tools\n",
        "\n",
        "\n",
        "def get_agent_async():\n",
        "    \"\"\"Asynchronously creates the MCP Toolset and the LlmAgent.\n",
        "\n",
        "    Returns:\n",
        "        tuple: (Agent instance)\n",
        "    \"\"\"\n",
        "    tools = get_tools_async()\n",
        "\n",
        "    root_agent = Agent(\n",
        "        model=\"gemini-2.5-pro\",  # Adjust model name if needed based on availability\n",
        "        name=\"hr_agent\",\n",
        "        instruction=agent_prompt,\n",
        "        tools=[tools],\n",
        "    )\n",
        "    print(\"LlmAgent created.\")\n",
        "\n",
        "    # Return both the agent and the exit_stack needed for cleanup\n",
        "    return root_agent\n",
        "\n",
        "\n",
        "global root_agent\n",
        "if root_agent is None:\n",
        "    log.info(\"Initializing agent...\")\n",
        "    root_agent = get_agent_async()\n",
        "    if root_agent:\n",
        "        log.info(\"Agent initialized successfully.\")\n",
        "    else:\n",
        "        log.error(\"Agent initialization failed.\")\n",
        "else:\n",
        "    log.info(\"Agent already initialized.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "be2c8b2153d3"
      },
      "source": [
        "##### 3.2 Evaluate ADK Agent Locally."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "abb9563867ee"
      },
      "outputs": [],
      "source": [
        "app = reasoning_engines.AdkApp(\n",
        "    agent=root_agent,\n",
        "    enable_tracing=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "600350db1ae3"
      },
      "outputs": [],
      "source": [
        "session = app.create_session(user_id=\"u_1231\")\n",
        "session"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5468bfe24dc3"
      },
      "outputs": [],
      "source": [
        "query = \"My employee id is 12 and want to apply for a leave from 21 July 2025 to 22 July 2025.\"\n",
        "contents = types.Content(role=\"user\", parts=[types.Part.from_text(text=query)])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d431f9f81184"
      },
      "outputs": [],
      "source": [
        "for event in app.stream_query(\n",
        "    user_id=\"u_1231\",\n",
        "    session_id=session.id,\n",
        "    message=contents.model_dump(),\n",
        "):\n",
        "    print(event)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "31bd98e3819f"
      },
      "source": [
        "#### Step 4: Agent Deployment to Agent Engine"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f4b19198f858"
      },
      "outputs": [],
      "source": [
        "from vertexai.preview.reasoning_engines import AdkApp\n",
        "\n",
        "app = AdkApp(agent=root_agent)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "308582742532"
      },
      "outputs": [],
      "source": [
        "pip install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c066be5562e3"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "PROJECT_NUMBER = PROJECT_NUMBER\n",
        "LOCATION = \"us-central1\"  # TODO for Developer : Update region here.\n",
        "STAGING_BUCKET = (\n",
        "    \"gs://[bucket-name]\"  # TODO for Developer : Update GCS bucket name here.\n",
        ")\n",
        "vertexai.init(project=PROJECT_NUMBER, location=LOCATION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7fded9eca51a"
      },
      "outputs": [],
      "source": [
        "from vertexai import agent_engines\n",
        "\n",
        "remote_app = agent_engines.create(\n",
        "    display_name=\"HR Agent V3\",\n",
        "    # agent_engine=root_agent,\n",
        "    agent_engine=app,\n",
        "    requirements=[\n",
        "        \"google-adk (==1.5.0)\",\n",
        "        \"google-genai (==1.24.0)\",\n",
        "        \"pydantic (==2.11.7)\",\n",
        "    ],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7bea20750e56"
      },
      "outputs": [],
      "source": [
        "remote_app.resource_name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9355b5b8f624"
      },
      "source": [
        "#### Step 5: Agent Engine Registration with Gemini Enterprise"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2b5a9e686af1"
      },
      "source": [
        "**Create oauth consent and mention the \"clientId\" and \"clientSecret\".**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3635a286d182"
      },
      "outputs": [],
      "source": [
        "%%bash\n",
        "\n",
        "curl -X POST \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
        "-H \"X-Goog-User-Project: [your-project-id]\" \\\n",
        "\"https://discoveryengine.googleapis.com/v1alpha/projects/[your-project-id]/locations/global/authorizations?authorizationId=customhr9893\" \\\n",
        "-d '{\n",
        "\"name\": \"projects/[your-project-id]/locations/global/authorizations/customhr9893\",\n",
        "\"serverSideOauth2\": {\n",
        "\"clientId\": \"[UPDATE-CLIENT-ID]\",\n",
        "\"clientSecret\": \"[UPDATE-CLIENT-SECRET]\",\n",
        "\"authorizationUri\": \"https://accounts.google.com/o/oauth2/v2/auth?client_id=[UPDATE-CLIENT-ID]&scope=https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fcloud-platform&include_granted_scopes=true&response_type=code&access_type=offline&prompt=consent\",\n",
        "\"tokenUri\": \"https://oauth2.googleapis.com/token\"\n",
        "}\n",
        "}'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b33096f67166"
      },
      "outputs": [],
      "source": [
        "%%bash\n",
        "\n",
        "curl -X POST \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
        "-H \"X-Goog-User-Project: [your-project-id]\" \\\n",
        "\"https://discoveryengine.googleapis.com/v1alpha/projects/[your-project-id]/locations/global/collections/default_collection/engines/[your-gemini-enterprise-engine-id]/assistants/default_assistant/agents\" \\\n",
        "-d '{\n",
        "\"displayName\": \"HR Agent V3\",\n",
        "\"description\": \"HR Agent can help employees to apply leave.\",\n",
        "\"adk_agent_definition\": {\n",
        "\"tool_settings\": {\n",
        "\"tool_description\": \"HR Agent can help employees to apply leave.\"\n",
        "},\n",
        "\"provisioned_reasoning_engine\": {\n",
        "\"reasoning_engine\": \"projects/[your-project-number]/locations/us-central1/reasoningEngines/[reasoning-engine-id]\"\n",
        "},\n",
        "\"authorizations\": [\"projects/[your-project-number]/locations/global/authorizations/customhr9893\"]\n",
        "}\n",
        "}'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "69085722b50a"
      },
      "source": [
        "#### After agent with MCP server is enabled, you can interact with agent on Gemini Enterprise. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "931d26753aca"
      },
      "source": [
        "![MCP_server_with_gemini_enterprise](https://services.google.com/fh/files/misc/mcp_with_agentspace_v2.png)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "mcp_with_gemini_enterprise.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
