{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eQheLG5T6YV3"
      },
      "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": "b4lYbaTN6YV4"
      },
      "source": [
        "# Get started with Cloud API Registry on Vertex AI Agent Engine\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/agents/agent_engine/tutorial_get_started_with_cloud_api_registry.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%2Fagents%2Fagent_engine%2Ftutorial_get_started_with_cloud_api_registry.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/agents/agent_engine/tutorial_get_started_with_cloud_api_registry.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/agents/agent_engine/tutorial_get_started_with_cloud_api_registry.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",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<p>\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/agents/agent_engine/tutorial_get_started_with_cloud_api_registry.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/agents/agent_engine/tutorial_get_started_with_cloud_api_registry.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/agents/agent_engine/tutorial_get_started_with_cloud_api_registry.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/agents/agent_engine/tutorial_get_started_with_cloud_api_registry.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/agents/agent_engine/tutorial_get_started_with_cloud_api_registry.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>\n",
        "</p>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tUD4rTEL6YV4"
      },
      "source": [
        "| Authors |\n",
        "| --- |\n",
        "| [Ivan Nardini](https://github.com/inardini) |\n",
        "| [Shawn Yang](https://github.com/shawn-yang-google) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QG69t3uX6YV5"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial shows you how to build and deploy an intelligent data analysis agent using **Cloud API Registry** with **ADK** and **Vertex AI Agent Engine**.\n",
        "\n",
        "You'll create an agent that can interact with BigQuery datasets through standardized MCP tools, allowing it to discover datasets, examine table schemas, and execute SQL queries—all through natural language conversations.\n",
        "\n",
        "### What you'll learn\n",
        "\n",
        "- How to discover and list available MCP servers and tools using Cloud API Registry\n",
        "- How to enable BigQuery MCP server tools for your project\n",
        "- How to build an agent using the Google Agent Development Kit (ADK) with Tool Registry\n",
        "- How to test your agent locally before deployment\n",
        "- How to deploy your agent to Vertex AI Agent Engine\n",
        "\n",
        "### What you'll build\n",
        "\n",
        "A data analyst agent that can:\n",
        "- List BigQuery datasets in your project\n",
        "- Retrieve table schemas and metadata\n",
        "- Execute SQL queries against public datasets\n",
        "- Provide natural language insights about your data\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "- A Google Cloud project with billing enabled\n",
        "- Basic familiarity with Python and BigQuery (helpful but not required)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L2Eiegvz6YV5"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZsjSIgrnqabk"
      },
      "source": [
        "### Install required packages\n",
        "\n",
        "**Note:** This installation will take a few moments. After it completes, you may see a message prompting you to restart the runtime. This is expected behavior—please click the **\"Restart Runtime\"** button when prompted."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w1MpoBuL6YV5"
      },
      "outputs": [],
      "source": [
        "# Install Vertex AI SDK with Agent Engine and ADK support\n",
        "%pip install \"google-cloud-aiplatform[agent_engines,adk]>=1.101.0\" --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GpJneV_T6YV6"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook in Google Colab, run the cell below to authenticate your Google Cloud account. This allows the notebook to access your Cloud resources."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b8yFBtZw6YV6"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()\n",
        "    print(\"✅ Authentication successful!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YUsoJeOz6YV6"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To use Vertex AI and Cloud API Registry, you need to specify your Google Cloud project details. Replace the placeholder values below with your actual project ID and preferred location.\n",
        "\n",
        "We'll also create a Cloud Storage bucket to store artifacts during agent deployment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "crHygQWb6YV6"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import vertexai\n",
        "\n",
        "# Replace with your project information\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "# fmt: on\n",
        "\n",
        "# Use environment variable if not explicitly set\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "# Set location (use us-central1 for best availability)\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "# Create a unique bucket name for this tutorial\n",
        "import uuid\n",
        "\n",
        "BUCKET_NAME = f\"{PROJECT_ID}-tool-registry-{uuid.uuid4().hex[:8]}\"\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "\n",
        "print(f\"📋 Project ID: {PROJECT_ID}\")\n",
        "print(f\"📍 Location: {LOCATION}\")\n",
        "print(f\"🪣 Bucket: {BUCKET_URI}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C8EI5nVr6YV6"
      },
      "source": [
        "### Create Cloud Storage bucket\n",
        "\n",
        "This bucket will store temporary files during agent deployment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "imEDRCAf6YV7"
      },
      "outputs": [],
      "source": [
        "# Create the bucket\n",
        "!gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}\n",
        "print(\"\\n✅ Bucket created successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KwKM8RHp6YV7"
      },
      "source": [
        "### Enable required Google Cloud APIs\n",
        "\n",
        "**Why we need these APIs:**\n",
        "- **apihub.googleapis.com**: Provides access to API Hub for discovering MCP servers\n",
        "- **cloudapiregistry.googleapis.com**: The Cloud API Registry service that manages MCP tools\n",
        "- **aiplatform.googleapis.com**: Vertex AI for deploying and running agents\n",
        "\n",
        "**Note:** This command may take 1-2 minutes to complete. This is normal."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cx43DGzm6YV7"
      },
      "outputs": [],
      "source": [
        "# Enable the required APIs\n",
        "!gcloud services enable apihub.googleapis.com cloudapiregistry.googleapis.com aiplatform.googleapis.com --project={PROJECT_ID}\n",
        "print(\"\\n✅ All required APIs enabled successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LFsyD8gT6YV7"
      },
      "source": [
        "### Initialize Vertex AI SDK and ADK\n",
        "\n",
        "Now we'll initialize the Vertex AI SDK and set environment variables required by the Agent Development Kit (ADK)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-qZWVjxH6YV7"
      },
      "outputs": [],
      "source": [
        "# Set ADK environment variables\n",
        "os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"TRUE\"\n",
        "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT_ID\n",
        "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = LOCATION\n",
        "\n",
        "# Initialize Vertex AI SDK\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)\n",
        "\n",
        "print(\"✅ Vertex AI and ADK initialized successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aPSW6Ttj6YV7"
      },
      "source": [
        "## Discover MCP Tools with Cloud API Registry\n",
        "\n",
        "Cloud API Registry acts as a centralized catalog of MCP (Model Context Protocol) servers. Each MCP server exposes a set of **tools** that agents can use to perform specific tasks.\n",
        "\n",
        "In this section, you'll:\n",
        "1. List all available MCP servers (both enabled and disabled)\n",
        "2. Filter for the BigQuery MCP server\n",
        "3. Examine what tools the BigQuery server provides\n",
        "\n",
        "Let's start by discovering what MCP servers are available in the registry."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2wmbKDcC6YV7"
      },
      "source": [
        "### List all available MCP servers\n",
        "\n",
        "The `--all` flag shows both enabled and disabled servers. Without it, you'd only see servers you've already enabled for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kwuJq0Fm6YV7"
      },
      "outputs": [],
      "source": [
        "# List all MCP servers available in the registry\n",
        "!gcloud beta api-registry mcp servers list --all --project={PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Clm16Mj-6YV7"
      },
      "source": [
        "### Find the BigQuery MCP server\n",
        "\n",
        "Now let's filter specifically for the BigQuery server by its URL. This shows us the server's status (enabled/disabled) and endpoint information."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I8GkZxKk6YV7"
      },
      "outputs": [],
      "source": [
        "# Filter for BigQuery MCP server\n",
        "!gcloud beta api-registry mcp servers list --all \\\n",
        "    --filter='urls=\"bigquery.googleapis.com/mcp\"' \\\n",
        "    --project={PROJECT_ID}\n",
        "\n",
        "print(\"💡 Note: If the state shows 'DISABLED', we'll enable it in the next section.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ro7XlXoD6YV7"
      },
      "source": [
        "### List available BigQuery MCP tools\n",
        "\n",
        "Each MCP server provides specific tools. Let's see what capabilities the BigQuery server offers.\n",
        "\n",
        "**Note:** This command lists tools even if the server isn't enabled yet. We're just exploring what's available."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fBiIslC76YV7"
      },
      "outputs": [],
      "source": [
        "# List all tools provided by the BigQuery MCP server\n",
        "!gcloud beta api-registry mcp tools list --all \\\n",
        "    --filter='mcpServerUrls=\"bigquery.googleapis.com/mcp\"' \\\n",
        "    --project={PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a-mVSWb46YV7"
      },
      "source": [
        "### Understanding BigQuery MCP Tools\n",
        "\n",
        "The BigQuery MCP server provides five powerful tools for working with your data:\n",
        "\n",
        "| Tool Name | Description | Use Case |\n",
        "|-----------|-------------|----------|\n",
        "| `list_dataset_ids` | Lists all dataset IDs in a project | Discovering what datasets are available |\n",
        "| `get_dataset_info` | Retrieves metadata about a specific dataset | Understanding dataset properties and settings |\n",
        "| `list_table_ids` | Lists all table IDs within a dataset | Exploring what tables exist in a dataset |\n",
        "| `get_table_info` | Retrieves schema and metadata for a table | Examining column names, types, and structure |\n",
        "| `execute_sql` | Runs a SQL query and returns results | Analyzing data with SELECT statements |\n",
        "\n",
        "**Important notes about `execute_sql`:**\n",
        "- ✅ Only `SELECT` statements are allowed\n",
        "- ❌ `INSERT`, `UPDATE`, `DELETE`, and stored procedures are blocked\n",
        "- All queries are automatically labeled with `goog-mcp-server: true` for tracking"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KadDhzjR6YV8"
      },
      "source": [
        "### Enable BigQuery MCP Server\n",
        "\n",
        "Before we can use the BigQuery tools in our agent, we need to enable the MCP server for our project. This grants our project permission to use the server's tools.\n",
        "\n",
        "**What happens when you enable a server:**\n",
        "- The server becomes available for agents in your project\n",
        "- Your project's agents can call the server's tools\n",
        "- The tools use your project's credentials and quotas"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iNtys-Hg6YV8"
      },
      "outputs": [],
      "source": [
        "# Enable the BigQuery MCP server\n",
        "!gcloud beta api-registry mcp enable bigquery.googleapis.com --project={PROJECT_ID}\n",
        "\n",
        "print(\"\\n✅ BigQuery MCP server enabled successfully!\")\n",
        "print(\"\\nYour agents can now use BigQuery tools to:\")\n",
        "print(\"  • List and explore datasets\")\n",
        "print(\"  • Examine table schemas\")\n",
        "print(\"  • Execute SQL queries\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5MqSkN4o6YV8"
      },
      "source": [
        "### Verify the server is enabled\n",
        "\n",
        "Let's confirm the BigQuery server now shows as `ENABLED` in our project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c7kttbph6YV8"
      },
      "outputs": [],
      "source": [
        "# List only enabled servers (should now include BigQuery)\n",
        "!gcloud beta api-registry mcp servers list --project={PROJECT_ID}\n",
        "\n",
        "print(\"\\n✅ The BigQuery server should now appear with state: ENABLED\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HnUsoMO26YV8"
      },
      "source": [
        "## Import required libraries\n",
        "\n",
        "Now we'll import the libraries needed to build our agent:\n",
        "- **ApiRegistry**: Connects to Cloud API Registry and retrieves MCP toolsets\n",
        "- **LlmAgent**: The core agent class from ADK\n",
        "- **Runner**: Allows us to test the agent locally\n",
        "- **InMemorySessionService**: Manages conversation history during local testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DSvrXpNe6YV8"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import uuid\n",
        "from collections.abc import Iterator\n",
        "from typing import Any\n",
        "\n",
        "from google.adk.agents import LlmAgent\n",
        "from google.adk.runners import Runner\n",
        "from google.adk.sessions import InMemorySessionService\n",
        "from google.adk.tools.api_registry import ApiRegistry\n",
        "from google.genai import types\n",
        "from vertexai import agent_engines\n",
        "from vertexai.agent_engines import AgentEngine\n",
        "from vertexai.preview.reasoning_engines import AdkApp\n",
        "\n",
        "print(\"✅ Libraries imported successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rTHvBPmg6YV8"
      },
      "source": [
        "## Build the Agent with Tool Registry\n",
        "\n",
        "Now comes the exciting part—building our data analyst agent!\n",
        "\n",
        "**How this works:**\n",
        "1. **ApiRegistry** connects to Cloud API Registry using your project ID\n",
        "2. We call `get_toolset()` to retrieve the BigQuery MCP tools\n",
        "3. We create an `LlmAgent` and equip it with these tools\n",
        "4. The agent can now use natural language to decide when and how to use each tool"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9DaH7pwuUT_u"
      },
      "source": [
        "### Grant IAM permissions for MCP Tool User role\n",
        "\n",
        "This allows your user account to call MCP tools through the API Registry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6RAqbL7qVLv-"
      },
      "outputs": [],
      "source": [
        "user_email = !gcloud config get-value account\n",
        "member = f\"user:{user_email[0]}\"\n",
        "\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member={member} --role=\"roles/mcp.toolUser\" --quiet --condition=None\n",
        "print(\"\\n✅ IAM permissions granted!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IYnwI91h6YV8"
      },
      "source": [
        "### Initialize the Tool Registry\n",
        "\n",
        "The `ApiRegistry` class authenticates with your Google Cloud project and fetches the list of enabled MCP servers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MFv_z7SC6YV8"
      },
      "outputs": [],
      "source": [
        "# Header required for BigQuery MCP server\n",
        "header_provider = lambda context: {\n",
        "    \"x-goog-user-project\": PROJECT_ID,\n",
        "}\n",
        "\n",
        "# Initialize the API Registry with your project\n",
        "tool_registry = ApiRegistry(PROJECT_ID, header_provider=header_provider)\n",
        "\n",
        "print(\"✅ Tool Registry initialized successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VBGvQGSI6YV8"
      },
      "source": [
        "### Get the BigQuery toolset\n",
        "\n",
        "Now we'll retrieve all the BigQuery tools. We could filter for specific tools using `tool_filter`, but we'll include all of them for maximum capability.\n",
        "\n",
        "**Behind the scenes:** This creates an authenticated connection to the BigQuery MCP server with proper headers and permissions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ReHzI_u6YV8"
      },
      "outputs": [],
      "source": [
        "# Get the BigQuery MCP toolset\n",
        "# Note: The server name follows the pattern \"projects/{PROJECT_ID}/locations/global/mcpServers/{service}\"\n",
        "registry_tools = tool_registry.get_toolset(\n",
        "    mcp_server_name=f\"projects/{PROJECT_ID}/locations/global/mcpServers/google-bigquery.googleapis.com-mcp\",\n",
        "    # Optional: filter for specific tools only\n",
        "    # tool_filter=[\"list_dataset_ids\", \"execute_sql\"]\n",
        ")\n",
        "\n",
        "print(\"✅ BigQuery toolset retrieved successfully!\")\n",
        "print(\"\\nThe agent now has access to all BigQuery MCP tools.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ANWGUb6-6YV8"
      },
      "source": [
        "### Create the agent\n",
        "\n",
        "Finally, we create our agent with:\n",
        "- **Model**: Gemini 2.0 Flash for fast, intelligent responses\n",
        "- **Name**: A descriptive identifier\n",
        "- **Instruction**: A clear directive that shapes the agent's behavior\n",
        "- **Tools**: The BigQuery toolset we just retrieved\n",
        "\n",
        "**The instruction is crucial** - it tells the agent how to behave and when to use its tools."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ECMxeeI56YV8"
      },
      "outputs": [],
      "source": [
        "# Create the LlmAgent with BigQuery tools\n",
        "data_analyst_agent = LlmAgent(\n",
        "    model=\"gemini-2.5-flash\",\n",
        "    name=\"bigquery_data_analyst\",\n",
        "    instruction=f\"\"\"\n",
        "    You are a helpful data analyst assistant with access to BigQuery.\n",
        "\n",
        "    The project ID is: {PROJECT_ID}\n",
        "\n",
        "    When users ask about data:\n",
        "    - Use the project ID {PROJECT_ID} when calling BigQuery tools\n",
        "    - First, explore available datasets and tables to understand what data exists\n",
        "    - Check table schemas to understand the structure before querying\n",
        "    - Write clear, efficient SQL queries to answer their questions\n",
        "    - Explain your findings in simple, non-technical language\n",
        "\n",
        "    Always use the BigQuery tools to fetch real data rather than making assumptions.\n",
        "    For all BigQuery operations, use project_id: {PROJECT_ID}\n",
        "    \"\"\",\n",
        "    tools=[registry_tools],\n",
        ")\n",
        "\n",
        "print(\"✅ Data analyst agent created successfully!\")\n",
        "print(\"\\n🤖 Agent capabilities:\")\n",
        "print(\"  • Natural language understanding\")\n",
        "print(\"  • BigQuery dataset exploration\")\n",
        "print(\"  • Schema inspection\")\n",
        "print(\"  • SQL query generation and execution\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9qZQeKGF6YV8"
      },
      "source": [
        "### Test the Agent Locally\n",
        "\n",
        "Before deploying to the cloud, it's important to test locally. This allows for faster iteration and easier debugging.\n",
        "\n",
        "**How local testing works:**\n",
        "1. **InMemorySessionService**: Stores conversation history in memory (no cloud storage needed)\n",
        "2. **Runner**: Orchestrates the conversation between you and the agent\n",
        "3. **Interactive loop**: You can chat with the agent to verify it works correctly\n",
        "\n",
        "Let's set up the local testing environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GgvV3Xlg6YWB"
      },
      "source": [
        "#### Create the interactive chat helper\n",
        "\n",
        "This helper function provides a consistent chat interface that works with both local and deployed agents. It handles:\n",
        "- Session management\n",
        "- Streaming responses\n",
        "- User input\n",
        "- Graceful exit"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A0wQO3Ky6YWB"
      },
      "outputs": [],
      "source": [
        "def chat_loop(app, user_id: str | None = None, session_id: str | None = None) -> None:\n",
        "    \"\"\"Interactive chat loop for testing agents.\n",
        "\n",
        "    Args:\n",
        "        app: Either a Runner (local) or AdkApp/AgentEngine (deployed)\n",
        "        user_id: Optional user identifier\n",
        "        session_id: Optional session identifier\n",
        "    \"\"\"\n",
        "    # Generate unique IDs if not provided\n",
        "    user_id = user_id or f\"user_{uuid.uuid4().hex[:8]}\"\n",
        "\n",
        "    # Handle session creation based on app type\n",
        "    if isinstance(app, (AdkApp, AgentEngine)):\n",
        "        # For deployed agents, create a remote session\n",
        "        if not session_id:\n",
        "            session = app.create_session(user_id=user_id)\n",
        "            session_id = session[\"id\"] if isinstance(session, dict) else session.id\n",
        "\n",
        "        def query_fn(msg: str):\n",
        "            return app.stream_query(user_id=user_id, session_id=session_id, message=msg)\n",
        "\n",
        "    elif isinstance(app, Runner):\n",
        "        # For local agents, use in-memory session\n",
        "        session_id = session_id or f\"session_{uuid.uuid4().hex[:8]}\"\n",
        "\n",
        "        def query_fn(msg: str):\n",
        "            return app.run(\n",
        "                user_id=user_id,\n",
        "                session_id=session_id,\n",
        "                new_message=types.Content(role=\"user\", parts=[types.Part(text=msg)]),\n",
        "            )\n",
        "\n",
        "    else:\n",
        "        raise TypeError(f\"Unsupported app type: {type(app)}\")\n",
        "\n",
        "    # Print session info\n",
        "    print(\"\\n🚀 Starting chat session...\")\n",
        "    print(f\"👤 User ID: {user_id}\")\n",
        "    print(f\"📁 Session ID: {session_id}\")\n",
        "    print(\"💬 Type 'exit' or 'quit' to end the session.\")\n",
        "    print(\"-\" * 60)\n",
        "\n",
        "    # Main chat loop\n",
        "    while True:\n",
        "        try:\n",
        "            user_input = input(\"\\nYou: \").strip()\n",
        "            if not user_input or user_input.lower() in {\"quit\", \"exit\", \"bye\"}:\n",
        "                break\n",
        "\n",
        "            print(\"\\nAssistant: \", end=\"\", flush=True)\n",
        "\n",
        "            # Stream and collect response (don't print again since _get_response_text already prints)\n",
        "            _get_response_text(query_fn(user_input))\n",
        "            print()  # Add newline after response\n",
        "\n",
        "        except (KeyboardInterrupt, EOFError):\n",
        "            print(\"\\n\\n🛑 Chat interrupted.\")\n",
        "            break\n",
        "        except Exception as e:\n",
        "            print(f\"\\n❌ Error: {e}\")\n",
        "            import traceback\n",
        "\n",
        "            traceback.print_exc()\n",
        "            continue\n",
        "\n",
        "    print(\"\\n👋 Goodbye!\")\n",
        "\n",
        "\n",
        "def _get_response_text(events: Iterator[Any]) -> None:\n",
        "    \"\"\"Extract and stream response text from events.\"\"\"\n",
        "    for event in events:\n",
        "        # Handle different event formats\n",
        "        if isinstance(event, dict):\n",
        "            text = _extract_from_dict_event(event)\n",
        "        else:\n",
        "            text = _extract_from_object_event(event)\n",
        "\n",
        "        if text:\n",
        "            print(text, end=\"\", flush=True)  # Stream to console\n",
        "\n",
        "\n",
        "def _extract_from_dict_event(event: dict) -> str | None:\n",
        "    \"\"\"Extract text from dictionary-style events.\"\"\"\n",
        "    # Skip function calls, only extract model text responses\n",
        "    content = event.get(\"content\", {})\n",
        "    if not content:\n",
        "        return None\n",
        "    if content.get(\"role\") == \"model\":\n",
        "        parts = content.get(\"parts\", [])\n",
        "        text_parts = []\n",
        "        for part in parts:\n",
        "            if isinstance(part, dict) and \"text\" in part:\n",
        "                text_parts.append(part[\"text\"])\n",
        "        return \"\".join(text_parts) if text_parts else None\n",
        "    return None\n",
        "\n",
        "\n",
        "def _extract_from_object_event(event: Any) -> str | None:\n",
        "    \"\"\"Extract text from object-style events.\"\"\"\n",
        "    content = getattr(event, \"content\", None)\n",
        "    if isinstance(content, str):\n",
        "        return content\n",
        "\n",
        "    if content and hasattr(content, \"parts\"):\n",
        "        text_parts = []\n",
        "        for part in content.parts:\n",
        "            if hasattr(part, \"text\") and part.text:\n",
        "                text_parts.append(part.text)\n",
        "        return \"\".join(text_parts) if text_parts else None\n",
        "\n",
        "    return getattr(event, \"text\", None)\n",
        "\n",
        "\n",
        "print(\"✅ Chat helper function defined!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-6iSuYPT6YWB"
      },
      "source": [
        "#### Initialize the local runner\n",
        "\n",
        "The Runner manages the conversation flow between you and the agent, handling session state and message history."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nW1XnOk36YWB"
      },
      "outputs": [],
      "source": [
        "# Create an in-memory session service for local testing\n",
        "session_service = InMemorySessionService()\n",
        "\n",
        "# Initialize the Runner with our agent\n",
        "runner = Runner(\n",
        "    agent=data_analyst_agent,\n",
        "    app_name=\"BigQueryDataAnalyst\",\n",
        "    session_service=session_service,\n",
        ")\n",
        "\n",
        "print(\"✅ Local runner initialized successfully!\")\n",
        "print(\"\\nYou're ready to test the agent locally.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lz8COm646YWB"
      },
      "source": [
        "#### Chat with your agent locally\n",
        "\n",
        "Now let's test the agent! Try asking questions like:\n",
        "\n",
        "**Example prompts:**\n",
        "- \"What datasets are available in my project?\"\n",
        "- \"Show me the tables in the bigquery-public-data.usa_names dataset\"\n",
        "- \"What's the schema of the usa_1910_current table?\"\n",
        "- \"Show me the top 10 most popular baby names in 2020\"\n",
        "\n",
        "**What to observe:**\n",
        "- The agent will use the BigQuery tools automatically\n",
        "- You'll see it make decisions about which tools to call\n",
        "- It will explain its findings in natural language\n",
        "\n",
        "Type 'exit' when you're done testing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d0IYomef_w_1"
      },
      "outputs": [],
      "source": [
        "async def test_agent_locally():\n",
        "    \"\"\"Test the agent locally with proper session creation.\"\"\"\n",
        "    # Create an in-memory session service for local testing\n",
        "    session_service = InMemorySessionService()\n",
        "\n",
        "    # Create a user ID\n",
        "    user_id = f\"user_{uuid.uuid4().hex[:8]}\"\n",
        "\n",
        "    # Create a session for this user\n",
        "    session = await session_service.create_session(\n",
        "        app_name=\"BigQueryDataAnalyst\", user_id=user_id\n",
        "    )\n",
        "\n",
        "    # Initialize the Runner with our agent\n",
        "    runner = Runner(\n",
        "        agent=data_analyst_agent,\n",
        "        app_name=\"BigQueryDataAnalyst\",\n",
        "        session_service=session_service,\n",
        "    )\n",
        "\n",
        "    print(\"✅ Local runner initialized successfully!\")\n",
        "    print(\"\\nYou're ready to test the agent locally.\")\n",
        "\n",
        "    # Start the interactive chat session with the created session\n",
        "    chat_loop(runner, user_id=user_id, session_id=session.id)\n",
        "\n",
        "\n",
        "# Run the async test function\n",
        "await test_agent_locally()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E-EZEOlW6YWB"
      },
      "source": [
        "## Deploy to Vertex AI Agent Engine\n",
        "\n",
        "Now that we've verified the agent works locally, let's deploy it to the cloud! This makes it:\n",
        "- Accessible via API\n",
        "- Scalable and production-ready\n",
        "- Shareable with your team\n",
        "\n",
        "**Deployment process:**\n",
        "1. Create a Python module file with the agent definition\n",
        "2. Use `agent_engines.create()` to deploy\n",
        "3. Agent Engine builds a container with your agent\n",
        "4. The agent becomes available as a managed service"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Srk4uxoc6YWB"
      },
      "source": [
        "### Create the agent module file\n",
        "\n",
        "**Why we need this:** The Tool Registry's `ApiRegistry` class contains non-serializable state, so we can't pickle the agent directly. Instead, we define it in a Python module that Agent Engine will import.\n",
        "\n",
        "**What this file does:**\n",
        "- Imports all necessary libraries\n",
        "- Reads configuration from environment variables\n",
        "- Defines a `VertexAiSessionService` builder for cloud-based sessions\n",
        "- Wraps the agent in an `AdkApp` for deployment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YwMiQMX-6YWB"
      },
      "outputs": [],
      "source": [
        "# Create the root_agent.py module file\n",
        "root_agent_file = f\"\"\"\n",
        "import os\n",
        "from google.adk.agents import LlmAgent\n",
        "from google.adk.tools.api_registry import ApiRegistry\n",
        "from vertexai.preview.reasoning_engines import AdkApp\n",
        "\n",
        "# Read configuration from environment variables\n",
        "# (These are set during deployment)\n",
        "PROJECT_ID = os.getenv('PROJECT_ID', '{PROJECT_ID}')\n",
        "LOCATION = os.getenv('LOCATION', '{LOCATION}')\n",
        "\n",
        "def session_service_builder():\n",
        "    '''Create a Vertex AI session service for cloud deployment.'''\n",
        "    from google.adk.sessions import VertexAiSessionService\n",
        "    return VertexAiSessionService(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "# Initialize Tool Registry and get BigQuery toolset\n",
        "header_provider = lambda context: {{\n",
        "    \"x-goog-user-project\": PROJECT_ID,\n",
        "}}\n",
        "\n",
        "# Initialize the API Registry with your project\n",
        "tool_registry = ApiRegistry(PROJECT_ID, header_provider=header_provider)\n",
        "registry_tools = tool_registry.get_toolset(\n",
        "    mcp_server_name=f\"projects/{{PROJECT_ID}}/locations/global/mcpServers/google-bigquery.googleapis.com-mcp\",\n",
        ")\n",
        "\n",
        "# Create the agent app\n",
        "agent_app = AdkApp(\n",
        "    agent=LlmAgent(\n",
        "        model='gemini-2.5-flash',\n",
        "        name='bigquery_data_analyst',\n",
        "        instruction=f'''You are a helpful data analyst assistant with access to BigQuery.\n",
        "\n",
        "    The project ID is: {{PROJECT_ID}}\n",
        "\n",
        "    When users ask about data:\n",
        "    - Use the project ID {{PROJECT_ID}} when calling BigQuery tools\n",
        "    - First, explore available datasets and tables to understand what data exists\n",
        "    - Check table schemas to understand the structure before querying\n",
        "    - Write clear, efficient SQL queries to answer their questions\n",
        "    - Explain your findings in simple, non-technical language\n",
        "\n",
        "    Always use the BigQuery tools to fetch real data rather than making assumptions.\n",
        "    For all BigQuery operations, use project_id: {{PROJECT_ID}}\n",
        "    ''',\n",
        "        tools=[registry_tools],\n",
        "    ),\n",
        "    session_service_builder=session_service_builder\n",
        ")\n",
        "\"\"\"\n",
        "\n",
        "# Write the file\n",
        "with open(\"root_agent.py\", \"w\") as f:\n",
        "    f.write(root_agent_file)\n",
        "\n",
        "print(\"✅ Agent module file created: root_agent.py\")\n",
        "print(\"\\nThis file will be uploaded to Agent Engine during deployment.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oE37Ix166YWB"
      },
      "source": [
        "### Deploy to Agent Engine\n",
        "\n",
        "Now we'll deploy the agent using `agent_engines.create()`. This process:\n",
        "- Uploads your agent module\n",
        "- Builds a Docker container with all dependencies\n",
        "- Deploys the container to a managed endpoint\n",
        "- Makes the agent available via API\n",
        "\n",
        "**Note:** This deployment typically takes 10-15 minutes. The progress will be displayed below.\n",
        "\n",
        "**What each parameter does:**\n",
        "- `display_name`: Human-readable name shown in the console\n",
        "- `agent_engine`: Specifies which module and variable contains your agent\n",
        "- `requirements`: Python packages to install in the container\n",
        "- `extra_packages`: Local files to include (your agent module)\n",
        "- `env_vars`: Environment variables available to your agent code"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Icvq6LgkD1co"
      },
      "source": [
        "#### Grant IAM permissions to the Vertex AI Agent Engine Service account\n",
        "\n",
        "This allows your Agent Engine account to call MCP tools through the API Registry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U75kzgEyD1cp"
      },
      "outputs": [],
      "source": [
        "# Get project number for Agent Engine service account\n",
        "project_number = !gcloud projects describe {PROJECT_ID} --format=\"value(projectNumber)\"\n",
        "agent_engine_sa = (\n",
        "    f\"service-{project_number[0]}@gcp-sa-aiplatform-re.iam.gserviceaccount.com\"\n",
        ")\n",
        "\n",
        "# Grant to Agent Engine service account\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member=serviceAccount:{agent_engine_sa} --role=\"roles/cloudapiregistry.viewer\" --quiet --condition=None\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member=serviceAccount:{agent_engine_sa} --role=\"roles/mcp.toolUser\" --quiet  --condition=None\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member=serviceAccount:{agent_engine_sa} --role=\"roles/bigquery.user\" --quiet --condition=None\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member=serviceAccount:{agent_engine_sa} --role=\"roles/bigquery.dataViewer\" --quiet --condition=None\n",
        "\n",
        "print(\"\\n✅ MCP Tool User role granted to Agent Engine service account!\")\n",
        "print(\"\\n💡 The deployed agent can now access MCP tools from API Registry\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GTc39UY9Ctkn"
      },
      "source": [
        "#### Create the startup script\n",
        "\n",
        "**Why we need this:** When deploying to Agent Engine, the BigQuery MCP server must be enabled in API Registry before the agent can use it.\n",
        "This script will run automatically during the build process to ensure the server is properly enabled.\n",
        "\n",
        "This setup step is crucial because Agent Engine runs in an isolated environment that needs explicit configuration of which MCP servers it can access."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TjBtqsgrCx4L"
      },
      "outputs": [],
      "source": [
        "# # Create startup_scripts directory if it doesn't exist\n",
        "os.makedirs(\"startup_scripts\", exist_ok=True)\n",
        "\n",
        "startup_script = f\"\"\"#!/bin/bash\n",
        "# installation_scripts/setup_api_registry.sh\n",
        "# This script enables the BigQuery MCP server in API Registry\n",
        "\n",
        "set -e  # Exit on error\n",
        "\n",
        "echo \"🔧 Start checking API Registry setup...\"\n",
        "\n",
        "# Get the project ID from environment variable\n",
        "PROJECT_ID=\"{PROJECT_ID}\"\n",
        "\n",
        "if [ -z \"$PROJECT_ID\" ]; then\n",
        "    echo \"❌ Error: PROJECT_ID environment variable not set\"\n",
        "    exit 1\n",
        "fi\n",
        "\n",
        "echo \"📋 Project ID: $PROJECT_ID\"\n",
        "\n",
        "# Install gcloud CLI if not already installed\n",
        "if ! command -v gcloud &> /dev/null; then\n",
        "    echo \"\"\n",
        "    echo \"📥 Installing Google Cloud SDK...\"\n",
        "\n",
        "    # Update package list and install dependencies\n",
        "    apt-get update -qq\n",
        "    apt-get install -y -qq curl apt-transport-https ca-certificates gnupg\n",
        "\n",
        "    # Add Google Cloud SDK repository\n",
        "    echo \"deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main\" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list\n",
        "\n",
        "    # Import Google Cloud public key\n",
        "    curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg\n",
        "\n",
        "    # Install gcloud CLI\n",
        "    apt-get update -qq\n",
        "    apt-get install -y -qq google-cloud-sdk google-cloud-sdk-gke-gcloud-auth-plugin\n",
        "\n",
        "    echo \"✅ Google Cloud SDK installed successfully!\"\n",
        "else\n",
        "    echo \"✅ Google Cloud SDK already installed\"\n",
        "fi\n",
        "\n",
        "# Verify the server is enabled\n",
        "echo \"\"\n",
        "echo \"🔍 Verifying server status...\"\n",
        "gcloud beta api-registry mcp servers list --project=\"$PROJECT_ID\"\n",
        "\n",
        "echo \"\"\n",
        "echo \"✅ API Registry setup complete!\"\n",
        "\"\"\"\n",
        "\n",
        "# Write the script\n",
        "with open(\"startup_scripts/check_api_registry.sh\", \"w\") as f:\n",
        "    f.write(startup_script)\n",
        "\n",
        "# Make it executable\n",
        "!chmod +x startup_scripts/check_api_registry.sh\n",
        "\n",
        "print(\"✅ Installation script created: startup_scripts/check_api_registry.sh\")\n",
        "print(\"\\n📦 This script will be included in the deployment package\")\n",
        "print(\"🚀 It will run automatically during the Agent Engine build process\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1J3JLgBtC8W-"
      },
      "source": [
        "#### Deploy the agent\n",
        "\n",
        "Now you are ready to deploy the agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yYhvR47w6YWB"
      },
      "outputs": [],
      "source": [
        "print(\"🚀 Starting deployment to Vertex AI Agent Engine...\")\n",
        "print(\"⏱️  This will take approximately 10-15 minutes.\")\n",
        "print(\"\\n\" + \"=\" * 60)\n",
        "\n",
        "# Deploy the agent\n",
        "remote_app = agent_engines.create(\n",
        "    display_name=\"bigquery-data-analyst\",\n",
        "    description=\"A data analyst agent with BigQuery access via Tool Registry\",\n",
        "    agent_engine=agent_engines.ModuleAgent(\n",
        "        module_name=\"root_agent\",  # The filename (without .py)\n",
        "        agent_name=\"agent_app\",  # The variable name in the module\n",
        "        register_operations={\n",
        "            \"\": [\"get_session\", \"list_sessions\", \"create_session\", \"delete_session\"],\n",
        "            \"async\": [\n",
        "                \"async_get_session\",\n",
        "                \"async_list_sessions\",\n",
        "                \"async_create_session\",\n",
        "                \"async_delete_session\",\n",
        "            ],\n",
        "            \"stream\": [\"stream_query\", \"streaming_agent_run_with_events\"],\n",
        "            \"async_stream\": [\"async_stream_query\"],\n",
        "        },\n",
        "    ),\n",
        "    requirements=[\n",
        "        \"google-cloud-aiplatform[agent_engines,adk]>=1.101.0\",\n",
        "    ],\n",
        "    extra_packages=[\n",
        "        \"root_agent.py\",\n",
        "        \"startup_scripts/check_api_registry.sh\",\n",
        "    ],\n",
        "    env_vars={\n",
        "        \"PROJECT_ID\": PROJECT_ID,\n",
        "        \"LOCATION\": LOCATION,\n",
        "    },\n",
        ")\n",
        "\n",
        "print(\"\\n\" + \"=\" * 60)\n",
        "print(\"✅ Deployment successful!\")\n",
        "print(\"\\n🎉 Your agent is now running in the cloud!\")\n",
        "print(\n",
        "    f\"\\n📊 View in console: https://console.cloud.google.com/vertex-ai/agents?project={PROJECT_ID}\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SGftRy7v6YWB"
      },
      "source": [
        "#### Test the Deployed Agent\n",
        "\n",
        "Your agent is now deployed and running in the cloud! Let's test it to ensure it works correctly.\n",
        "\n",
        "**What's different from local testing:**\n",
        "- Sessions are stored in Vertex AI (persist across runs)\n",
        "- The agent runs on managed infrastructure\n",
        "- You're accessing it via API calls\n",
        "\n",
        "**Try the same prompts you used locally:**\n",
        "- \"What datasets are available?\"\n",
        "- \"Show me tables in bigquery-public-data.usa_names\"\n",
        "- \"What are the top 10 baby names from 2020?\"\n",
        "\n",
        "Type 'exit' when you're done testing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wuXNU3_d6YWC"
      },
      "outputs": [],
      "source": [
        "# Start an interactive chat session with the deployed agent\n",
        "chat_loop(remote_app)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N5HzbvJl6YWC"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To avoid ongoing charges to your Google Cloud account, delete the resources you created in this tutorial.\n",
        "\n",
        "**Resources to clean up:**\n",
        "- The deployed Agent Engine instance\n",
        "- The Cloud Storage bucket\n",
        "\n",
        "**Note:** If you want to keep the agent for later use, skip the cleanup section."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1tYiJX4v6YWC"
      },
      "outputs": [],
      "source": [
        "# Set these flags to control what gets deleted\n",
        "delete_agent_engine = True  # @param {type:\"boolean\"}\n",
        "delete_bucket = True  # @param {type:\"boolean\"}\n",
        "\n",
        "if delete_agent_engine:\n",
        "    print(\"🗑️  Deleting Agent Engine instance...\")\n",
        "    # Find and delete all agents with our display name\n",
        "    agents = agent_engines.list(filter=\"display_name=bigquery-data-analyst\")\n",
        "    for agent in agents:\n",
        "        agent.delete(force=True)\n",
        "    print(\"✅ Agent Engine instance deleted successfully!\")\n",
        "\n",
        "if delete_bucket:\n",
        "    print(\"\\n🗑️  Deleting Cloud Storage bucket...\")\n",
        "    !gsutil rm -r {BUCKET_URI}\n",
        "    print(\"✅ Bucket deleted successfully!\")\n",
        "\n",
        "print(\"\\n✅ Cleanup complete!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hSxaS5mZ6YWC"
      },
      "source": [
        "## Next Steps\n",
        "\n",
        "Congratulations! You've successfully:\n",
        "- Discovered MCP servers using Cloud API Registry\n",
        "- Enabled BigQuery MCP tools\n",
        "- Built a data analyst agent with Tool Registry\n",
        "- Tested the agent locally\n",
        "- Deployed it to Vertex AI Agent Engine\n",
        "\n",
        "### What to explore next:\n",
        "\n",
        "1. **Customize the agent**: Modify the instruction to specialize in specific types of analysis\n",
        "2. **Filter tools**: Use `tool_filter` to limit which BigQuery tools the agent can access\n",
        "3. **Add more MCP servers**: Explore other MCP servers in the registry\n",
        "4. **Build multi-agent systems**: Combine multiple agents with different capabilities\n",
        "5. **Integrate with applications**: Use the Agent Engine API to embed your agent in applications\n",
        "\n",
        "### Learn more:\n",
        "\n",
        "- [Cloud API Registry documentation](https://docs.cloud.google.com/api-registry/docs/overview)\n",
        "- [Vertex AI Agent Engine documentation](cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview)\n",
        "- [Google Agent Development Kit (ADK)](https://google.github.io/adk-docs/)\n",
        "- [Model Context Protocol](https://modelcontextprotocol.io/)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tutorial_get_started_with_cloud_api_registry.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
