{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "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": "a9f992c63ce7"
      },
      "source": [
        "# Develop Model Context Protocol (MCP) Client and Server with Google Agent Development Kit (ADK)\n",
        "\n",
        "\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/mcp/develop_mcp_with_gemini_and_adk.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%2Fmcp%2Fdevelop_mcp_with_gemini_and_adk.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/mcp/develop_mcp_with_gemini_and_adk.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",
        "  \n",
        "  \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/mcp/develop_mcp_with_gemini_and_adk.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",
        "<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/mcp/develop_mcp_with_gemini_and_adk.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/mcp/develop_mcp_with_gemini_and_adk.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/mcp/develop_mcp_with_gemini_and_adk.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/mcp/develop_mcp_with_gemini_and_adk.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/mcp/develop_mcp_with_gemini_and_adk.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": "84f0f73a0f76"
      },
      "source": [
        "| Author(s) |\n",
        "| --- |\n",
        "| [Dave Wang](https://github.com/wadave) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Agent Development Kit (ADK) is a flexible and modular framework for developing and deploying AI agents. ADK can be used with popular LLMs and open-source generative AI tools and is designed with a focus on tight integration with the Google ecosystem and Gemini models. ADK makes it easy to get started with simple agents with Gemini models and Google AI tools while providing the control and structure needed for more complex agent architectures and orchestration. \n",
        "\n",
        "The Model Context Protocol (MCP) is an open standard that simplifies how AI assistants connect with external data, tools, and systems. It achieves this by standardizing the way applications provide contextual information to Large Language Models (LLMs), creating a vital interface for models to interact directly with various external services.\n",
        "\n",
        "Developers building MCP-enabled applications have the flexibility to utilize existing third-party MCP servers or implement their own custom server solutions.\n",
        "\n",
        "This notebook focuses on the latter, demonstrating how to build custom MCP servers using Gemini. Then we will show to how to use ADK with MCP clients to communicate with MCP servers.\n",
        "\n",
        "#### MCP server code generation examples:\n",
        "- Building a Cocktail MCP Server\n",
        "\n",
        "#### ADK and MCP client integration:\n",
        "-  Use Google ADK single agent to test\n",
        "-  Use Google ADK multi-agent to test"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Google Gen AI SDK and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-genai mcp geopy black google-cloud-bigquery google-adk"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dd2d67b1696a"
      },
      "source": [
        "### Import Libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f4b19d8bc07a"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "from pathlib import Path\n",
        "import re\n",
        "import sys\n",
        "\n",
        "import black\n",
        "from google import genai\n",
        "from google.genai import types\n",
        "from google.genai.types import (\n",
        "    GenerateContentConfig,\n",
        ")\n",
        "from mcp import ClientSession, StdioServerParameters\n",
        "from mcp.client.stdio import stdio_client\n",
        "import requests"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a18710bb1d4f"
      },
      "source": [
        "### Helper function\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ccda55e3eaf0"
      },
      "outputs": [],
      "source": [
        "def get_url_content(url):\n",
        "    try:\n",
        "        # Send an HTTP GET request to the URL\n",
        "        response = requests.get(url)\n",
        "\n",
        "        # Raise an exception if the request returned an error status code (like 404 or 500)\n",
        "        response.raise_for_status()\n",
        "\n",
        "        # Get the content of the response as text (HTML, in this case)\n",
        "        # 'requests' automatically decodes the content based on HTTP headers\n",
        "        file_content = response.text\n",
        "\n",
        "        # Now you can work with the content\n",
        "        print(\"Successfully fetched content\")\n",
        "        return file_content\n",
        "        # Or, save it to a file:\n",
        "        # with open(\"server_page.html\", \"w\", encoding=\"utf-8\") as f:\n",
        "        #     f.write(file_content)\n",
        "        # print(\"Content saved to server_page.html\")\n",
        "\n",
        "    except requests.exceptions.RequestException as e:\n",
        "        # Handle potential errors during the request (e.g., network issues, DNS errors)\n",
        "        print(f\"Error fetching URL {url}: {e}\")\n",
        "    except requests.exceptions.HTTPError as e:\n",
        "        # Handle HTTP error responses (e.g., 404 Not Found, 503 Service Unavailable)\n",
        "        print(f\"HTTP Error for {url}: {e}\")\n",
        "\n",
        "\n",
        "def format_python(raw_code, output_filename):\n",
        "\n",
        "    try:\n",
        "        # Format the code string using black\n",
        "        # Use default FileMode which is generally recommended\n",
        "        formatted_code = black.format_str(raw_code, mode=black.FileMode())\n",
        "\n",
        "        # Save the formatted code to the specified file\n",
        "        with open(output_filename, \"w\", encoding=\"utf-8\") as f:\n",
        "            f.write(formatted_code)\n",
        "\n",
        "        print(f\"Successfully formatted the code and saved it to '{output_filename}'\")\n",
        "\n",
        "    except black.InvalidInput as e:\n",
        "        print(\n",
        "            \"Error formatting code: The input string does not seem to be valid Python syntax.\"\n",
        "        )\n",
        "        print(f\"Details: {e}\")\n",
        "    except Exception as e:\n",
        "        print(f\"An error occurred while writing the file: {e}\")\n",
        "\n",
        "\n",
        "def extract_json_from_string(input_str: str) -> dict | list | None:\n",
        "    \"\"\"\n",
        "    Extracts JSON data from a string, handling potential variations.\n",
        "\n",
        "    This function attempts to find JSON data within a string. It specifically\n",
        "    looks for JSON enclosed in Markdown-like code fences (```json ... ```).\n",
        "    If such a block is found, it extracts and parses the content.\n",
        "    If no code block is found, it attempts to parse the entire input string\n",
        "    as JSON.\n",
        "\n",
        "    Args:\n",
        "        input_str: The string potentially containing JSON data. It might be\n",
        "                   a plain JSON string or contain a Markdown code block\n",
        "                   with JSON, possibly preceded by other text (like 'shame').\n",
        "\n",
        "    Returns:\n",
        "        The parsed JSON object (typically a dictionary or list) if valid\n",
        "        JSON is found and successfully parsed.\n",
        "        Returns None if no valid JSON is found, if parsing fails, or if the\n",
        "        input is not a string.\n",
        "    \"\"\"\n",
        "    if not isinstance(input_str, str):\n",
        "        # Handle cases where input is not a string\n",
        "        return None\n",
        "\n",
        "    # Pattern to find JSON within ```json ... ``` blocks\n",
        "    # - ````json`: Matches the start fence.\n",
        "    # - `\\s*`: Matches any leading whitespace after the fence marker.\n",
        "    # - `(.*?)`: Captures the content (non-greedily) between the fences. This is group 1.\n",
        "    # - `\\s*`: Matches any trailing whitespace before the end fence.\n",
        "    # - ` ``` `: Matches the end fence.\n",
        "    # - `re.DOTALL`: Allows '.' to match newline characters.\n",
        "    pattern = r\"```json\\s*(.*?)\\s*```\"\n",
        "    match = re.search(pattern, input_str, re.DOTALL)\n",
        "\n",
        "    json_string_to_parse = None\n",
        "\n",
        "    if match:\n",
        "        # If a markdown block is found, extract its content\n",
        "        json_string_to_parse = match.group(\n",
        "            1\n",
        "        ).strip()  # Get captured group and remove surrounding whitespace\n",
        "    else:\n",
        "        # If no markdown block, assume the *entire* input might be JSON\n",
        "        # We strip whitespace in case the string is just JSON with padding\n",
        "        json_string_to_parse = input_str.strip()\n",
        "\n",
        "    if not json_string_to_parse:\n",
        "        # If after stripping, the potential JSON string is empty, return None\n",
        "        return None\n",
        "\n",
        "    try:\n",
        "        # Attempt to parse the determined string (either from block or whole input)\n",
        "        parsed_json = json.loads(json_string_to_parse)\n",
        "        return parsed_json\n",
        "    except json.JSONDecodeError:\n",
        "        # Parsing failed, indicating the string wasn't valid JSON\n",
        "        return None\n",
        "    except Exception as e:\n",
        "        # Catch other potential unexpected errors during parsing\n",
        "        print(f\"An unexpected error occurred during JSON parsing: {e}\")\n",
        "        return None\n",
        "\n",
        "\n",
        "def create_folder_if_not_exists(folder_path_str: str) -> bool:\n",
        "    \"\"\"\n",
        "    Creates a folder (and any necessary parent folders) if it doesn't already exist.\n",
        "    Uses print() for status and error messages.\n",
        "\n",
        "    Args:\n",
        "        folder_path_str (str): The path string for the folder to be created.\n",
        "                               Can be relative or absolute.\n",
        "\n",
        "    Returns:\n",
        "        bool: True if the folder already exists or was successfully created,\n",
        "              False if an error occurred during creation (e.g., permission denied).\n",
        "    \"\"\"\n",
        "    try:\n",
        "        # Convert the string path to a Path object\n",
        "        folder_path = Path(folder_path_str)\n",
        "\n",
        "        # Use mkdir() with options:\n",
        "        # parents=True: Creates any necessary parent directories. Like 'mkdir -p'.\n",
        "        # exist_ok=True: Doesn't raise an error if the directory already exists.\n",
        "        folder_path.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "        # Print confirmation (using resolve() to show the absolute path)\n",
        "        print(f\"Info: Successfully ensured folder exists: {folder_path.resolve()}\")\n",
        "        return True\n",
        "\n",
        "    except PermissionError:\n",
        "        print(\n",
        "            f\"Error: Permission denied: Could not create folder at '{folder_path_str}'.\"\n",
        "        )\n",
        "        return False\n",
        "    except OSError as e:\n",
        "        # Catch other OS-related errors (e.g., path is a file, invalid path format on Windows)\n",
        "        print(f\"Error: OS error creating folder '{folder_path_str}': {e}\")\n",
        "        return False\n",
        "    except Exception as e:\n",
        "        # Catch any other unexpected errors\n",
        "        print(\n",
        "            f\"Error: An unexpected error occurred creating folder '{folder_path_str}': {e}\"\n",
        "        )\n",
        "        return False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "128873119c0d"
      },
      "outputs": [],
      "source": [
        "create_folder_if_not_exists(\"server\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Option 1 use a  Vertex AI project\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a63e471cbadd"
      },
      "source": [
        "### Option 2. Use a Google API Key \n",
        "Uncomment the following block to use Express Mode"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a551b3691b32"
      },
      "outputs": [],
      "source": [
        "# API_KEY = \"[your-api-key]\"  # @param {type: \"string\", placeholder: \"[your-api-key]\", isTemplate: true}\n",
        "\n",
        "# if not API_KEY or API_KEY == \"[your-api-key]\":\n",
        "#     raise Exception(\"You must provide an API key to use Google AI in express mode.\")\n",
        "\n",
        "# client = genai.Client(api_key=API_KEY)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eee1d4f563fe"
      },
      "source": [
        "## Set up model id"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "259a8b9e467b"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"gemini-2.5-pro-preview-03-25\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3b68c286f758"
      },
      "source": [
        "### Get system instruction context info"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9ee9559094ab"
      },
      "outputs": [],
      "source": [
        "# The URL you want to fetch\n",
        "url = \"https://modelcontextprotocol.io/quickstart/server\"\n",
        "reference_content = get_url_content(url)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "369c0049239e"
      },
      "source": [
        "### Set up system instruction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d9df35a9d385"
      },
      "outputs": [],
      "source": [
        "from pydantic import BaseModel\n",
        "\n",
        "\n",
        "class ResponseSchema(BaseModel):\n",
        "    python_code: str\n",
        "    description: str\n",
        "\n",
        "\n",
        "system_instruction = f\"\"\"\n",
        "  You are an MCP server export.\n",
        "  Your mission is to write python code for MCP server.\n",
        "  Here's the MCP server development guide and example\n",
        "  {reference_content}\n",
        "  \n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "49ffc3ba5165"
      },
      "source": [
        "#### Set function to generate MCP server code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "67c8556ef889"
      },
      "outputs": [],
      "source": [
        "def generate_mcp_server(prompt):\n",
        "    response = client.models.generate_content(\n",
        "        model=MODEL_ID,\n",
        "        contents=prompt,\n",
        "        config=GenerateContentConfig(\n",
        "            system_instruction=system_instruction,\n",
        "            response_mime_type=\"application/json\",\n",
        "            response_schema=ResponseSchema,\n",
        "        ),\n",
        "    )\n",
        "\n",
        "    return response.text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "98e7d53203b9"
      },
      "source": [
        "## Generate MCP Server Code"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c7189e4c572c"
      },
      "source": [
        "### Build MCP Server for the Cocktail DB"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5c147f6f14bc"
      },
      "outputs": [],
      "source": [
        "ct_url = \"https://www.thecocktaildb.com/api.php\"\n",
        "ct_prompt_base = \"\"\"\n",
        "  Please create an MCP server code for the cocktail db. It has 5 tools:\n",
        "  1. search cocktail by name\n",
        "  2. list all cocktail by first letter\n",
        "  3. search ingredient by name. \n",
        "  4. list random cocktails\n",
        "  5. lookup full cocktail details by id\n",
        "  \n",
        "  Here's the API details:\n",
        "\n",
        "\"\"\"\n",
        "prompt = [ct_prompt_base, types.Part.from_uri(file_uri=ct_url, mime_type=\"text/html\")]\n",
        "response_text = generate_mcp_server(prompt)\n",
        "python_code = extract_json_from_string(response_text)[\"python_code\"]\n",
        "\n",
        "format_python(python_code, \"server/cocktail.py\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a7425e952e79"
      },
      "source": [
        "## Testing MCP Servers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8dffc0be1a8b"
      },
      "source": [
        "### Testing with Google ADK\n",
        "Note: It may not work in Colab (as of 4/16/2025) due to std io limitations. Please run it in Jupyter Notebook (VSCode, or Vertex AI workbench). "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b9935caa8148"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"1\"\n",
        "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT_ID\n",
        "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = LOCATION"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "421aaee02c2d"
      },
      "outputs": [],
      "source": [
        "import contextlib\n",
        "\n",
        "from dotenv import load_dotenv\n",
        "from google.adk.agents.llm_agent import LlmAgent\n",
        "from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService\n",
        "from google.adk.runners import Runner\n",
        "from google.adk.sessions import InMemorySessionService\n",
        "from google.adk.tools.mcp_tool.mcp_toolset import (\n",
        "    MCPToolset,\n",
        "    StdioServerParameters,\n",
        ")\n",
        "from google.genai import types\n",
        "\n",
        "load_dotenv()\n",
        "\n",
        "\n",
        "async def get_tools_async(server_params):\n",
        "    \"\"\"Gets tools from MCP Server.\"\"\"\n",
        "    tools, exit_stack = await MCPToolset.from_server(connection_params=server_params)\n",
        "    # MCP requires maintaining a connection to the local MCP Server.\n",
        "    # Using exit_stack to clean up server connection before exit.\n",
        "    return tools, exit_stack\n",
        "\n",
        "\n",
        "async def get_agent_async(server_params):\n",
        "    \"\"\"Creates an ADK Agent with tools from MCP Server.\"\"\"\n",
        "    tools, exit_stack = await get_tools_async(server_params)\n",
        "    root_agent = LlmAgent(\n",
        "        model=\"gemini-2.5-pro-preview-03-25\",\n",
        "        name=\"ai_assistant\",\n",
        "        instruction=\"Use tools to get information to answer user questions\",\n",
        "        tools=tools,\n",
        "    )\n",
        "    return root_agent, exit_stack"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "25805241904d"
      },
      "outputs": [],
      "source": [
        "async def list_mcp_tools(server_params):\n",
        "    async with stdio_client(server_params) as (read, write):\n",
        "        async with ClientSession(\n",
        "            read,\n",
        "            write,\n",
        "        ) as session:\n",
        "\n",
        "            # Initialize the connection\n",
        "            await session.initialize()\n",
        "\n",
        "            # Get tools from MCP session and convert to Gemini Tool objects\n",
        "            mcp_tools = await session.list_tools()\n",
        "\n",
        "            return mcp_tools"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c43e4944992"
      },
      "source": [
        "### A Google ADK agent integrated with single MCP client"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "958bb89b74e4"
      },
      "outputs": [],
      "source": [
        "async def run_adk_agent(server_params, question):\n",
        "    session_service = InMemorySessionService()\n",
        "    artifacts_service = InMemoryArtifactService()\n",
        "    session = session_service.create_session(state={}, app_name=\"my_app\", user_id=\"123\")\n",
        "\n",
        "    query = question\n",
        "    print(\"[user]: \", query)\n",
        "    content = types.Content(role=\"user\", parts=[types.Part(text=query)])\n",
        "    root_agent, exit_stack = await get_agent_async(server_params)\n",
        "    runner = Runner(\n",
        "        app_name=\"my_app\",\n",
        "        agent=root_agent,\n",
        "        artifact_service=artifacts_service,\n",
        "        session_service=session_service,\n",
        "    )\n",
        "    events_async = runner.run_async(\n",
        "        session_id=session.id, user_id=\"123\", new_message=content\n",
        "    )\n",
        "\n",
        "    async for event in events_async:\n",
        "        # print(event)\n",
        "        if event.content.role == \"user\" and event.content.parts[0].text:\n",
        "            print(\"[user]:\", event.content.parts[0].text)\n",
        "        if event.content.parts[0].function_response:\n",
        "            print(\"[-tool_response-]\", event.content.parts[0].function_response)\n",
        "        if event.content.role == \"model\" and event.content.parts[0].text:\n",
        "            print(\"[agent]:\", event.content.parts[0].text)\n",
        "\n",
        "    await exit_stack.aclose()\n",
        "    return events_async"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "31e57a5820b0"
      },
      "outputs": [],
      "source": [
        "ct_server_params = StdioServerParameters(\n",
        "    command=\"python\",\n",
        "    args=[\"./server/cocktail2.py\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ec87f962791d"
      },
      "outputs": [],
      "source": [
        "events_async = await run_adk_agent(\n",
        "    ct_server_params,\n",
        "    \"Please get cocktail margarita id and then full detail of cocktail margarita\",\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a4a1fc6223b5"
      },
      "outputs": [],
      "source": [
        "bq_server_params = StdioServerParameters(\n",
        "    command=\"python\",\n",
        "    # Make sure to update to the full absolute path to your server file\n",
        "    args=[\"./server/bq.py\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d50cbf4b8745"
      },
      "outputs": [],
      "source": [
        "await run_adk_agent(\n",
        "    bq_server_params,\n",
        "    \"Please list my BigQuery tables, project id is 'dw-genai-dev', location is 'us'\",\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "55cff367fd93"
      },
      "outputs": [],
      "source": [
        "med_server_params = StdioServerParameters(\n",
        "    command=\"python\",\n",
        "    # Make sure to update to the full absolute path to your server file\n",
        "    args=[\"./server/med.py\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1b7a33e7e189"
      },
      "outputs": [],
      "source": [
        "await run_adk_agent(med_server_params, \"Please explain flu in detail.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d5e1fe09e204"
      },
      "outputs": [],
      "source": [
        "nih_server_params = StdioServerParameters(\n",
        "    command=\"python\",\n",
        "    # Make sure to update to the full absolute path to your server file\n",
        "    args=[\"./server/nih.py\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "946b89f7db31"
      },
      "outputs": [],
      "source": [
        "await run_adk_agent(nih_server_params, \"Please tell me icd-10 code for pneumonia\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aad3e772d932"
      },
      "outputs": [],
      "source": [
        "ct_server_params = StdioServerParameters(\n",
        "    command=\"python\",\n",
        "    args=[\"./server/cocktail2.py\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b8b83eb3a429"
      },
      "outputs": [],
      "source": [
        "await run_adk_agent(\n",
        "    ct_server_params,\n",
        "    \"Please get cocktail margarita id and then full detail of cocktail margarita\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61c275e388f3"
      },
      "source": [
        "### Test ADK multi-agent with MCP multi-client"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0f17ea9a4542"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"gemini-2.0-flash\"\n",
        "from pydantic import BaseModel\n",
        "\n",
        "\n",
        "class AllServerConfigs(BaseModel):\n",
        "    configs: dict[str, StdioServerParameters]\n",
        "\n",
        "\n",
        "async def run_multi_agent_with_mcp_clients(\n",
        "    server_config_dict: AllServerConfigs, query: str\n",
        "):\n",
        "    session_service = InMemorySessionService()\n",
        "    artifacts_service = InMemoryArtifactService()\n",
        "    session = session_service.create_session(state={}, app_name=\"my_app\", user_id=\"123\")\n",
        "\n",
        "    print(\"[user]: \", query)\n",
        "    content = types.Content(role=\"user\", parts=[types.Part(text=query)])\n",
        "\n",
        "    all_tools = {}\n",
        "    # Use a single ExitStack in the main task\n",
        "    async with contextlib.AsyncExitStack() as stack:  # main stack\n",
        "        print(\"Setting up MCP connections sequentially...\")\n",
        "        for key, value in server_config_dict.items():\n",
        "            server_params = value\n",
        "            individual_exit_stack = (\n",
        "                None  # Define outside try for broader scope if needed\n",
        "            )\n",
        "            try:\n",
        "                # 1. AWAIT the call to run the function and get its results\n",
        "                print(f\"  Attempting connection for {server_params}...\")\n",
        "                tools, individual_exit_stack = await MCPToolset.from_server(\n",
        "                    connection_params=server_params\n",
        "                )\n",
        "\n",
        "                # 2. Check if an exit stack was actually returned\n",
        "                if individual_exit_stack is None:\n",
        "                    print(\n",
        "                        f\"  Warning: No exit stack returned for {server_params}. Cannot manage cleanup.\"\n",
        "                    )\n",
        "\n",
        "                # 3. Enter the individual_exit_stack into the main stack\n",
        "                #    This makes the main stack responsible for cleaning it up later.\n",
        "                print(f\"  Registering cleanup stack for {server_params}...\")\n",
        "                await stack.enter_async_context(individual_exit_stack)\n",
        "\n",
        "                # 4. Add the tools\n",
        "                print(f\"  Connection established for {server_params}, got tools.\")\n",
        "                # Check if tools is None or empty if connection might partially fail\n",
        "                if tools:\n",
        "                    all_tools.update({key: tools})\n",
        "                else:\n",
        "                    print(\n",
        "                        f\"  Warning: Connection successful but no tools returned for {server_params}.\"\n",
        "                    )\n",
        "\n",
        "            except TypeError as te:\n",
        "                print(f\"TypeError during setup for {server_params}: {te}\")\n",
        "\n",
        "                # Decide whether to continue or raise\n",
        "            except Exception as e:\n",
        "                print(f\"Error setting up connection for {server_params}: {e}\")\n",
        "\n",
        "        print(f\"Finished setup. Collected {len(all_tools)} servers.\")\n",
        "\n",
        "        # --- Agent Creation and Run (remains the same) ---\n",
        "        if not all_tools:\n",
        "            print(\n",
        "                \"Warning: No tools were collected. Agent may not function as expected.\"\n",
        "            )\n",
        "\n",
        "        print(all_tools)\n",
        "        booking_tools = all_tools[\"bnb\"]\n",
        "        booking_tools.extend(all_tools[\"weather\"])\n",
        "\n",
        "        ct_tools = all_tools[\"ct\"]\n",
        "\n",
        "        booking_agent = LlmAgent(\n",
        "            model=MODEL_ID,\n",
        "            name=\"booking_assistant\",\n",
        "            instruction=\"Use tools to get information to answer user questions\",\n",
        "            tools=booking_tools,\n",
        "        )\n",
        "\n",
        "        cocktail_agent = LlmAgent(\n",
        "            model=MODEL_ID,\n",
        "            name=\"cocktail_assistant\",\n",
        "            instruction=\"Use tools to get information to answer user questions\",\n",
        "            tools=ct_tools,\n",
        "        )\n",
        "\n",
        "        root_agent = LlmAgent(\n",
        "            model=MODEL_ID,\n",
        "            name=\"ai_assistant\",\n",
        "            instruction=\"\"\"You have access to sub-agents named 'cocktail_assistant' and 'booking_assistant. \n",
        "            - If the user asks about cocktails, delegate the task\n",
        "            to the 'cocktail_assistant' sub-agent. \n",
        "            - If the user asks about weather, room or house booking, delegate the task\n",
        "            to the 'booking_assistant' sub-agent.\n",
        "            - Carefully combine the information you find into a complete answer.\n",
        "            - If you cannot find the specific information requested using your tools, let the user know.\n",
        "            - Please format your response using Markdown to make it easy to read and understand.\n",
        "            \"\"\",\n",
        "            sub_agents=[cocktail_agent, booking_agent],\n",
        "        )\n",
        "\n",
        "        runner = Runner(\n",
        "            app_name=\"my_app\",\n",
        "            agent=root_agent,\n",
        "            artifact_service=artifacts_service,\n",
        "            session_service=session_service,\n",
        "        )\n",
        "\n",
        "        print(\"Running agent...\")\n",
        "        events_async = runner.run_async(\n",
        "            session_id=session.id, user_id=\"123\", new_message=content\n",
        "        )\n",
        "\n",
        "        async for event in events_async:\n",
        "            # Your event processing logic...\n",
        "            if event.content.role == \"user\" and event.content.parts[0].text:\n",
        "                print(\"[user]:\", event.content.parts[0].text)\n",
        "            if event.content.parts[0].function_response:\n",
        "                print(\"[-tool_response-]\", event.content.parts[0].function_response)\n",
        "            if event.content.role == \"model\" and event.content.parts[0].text:\n",
        "                print(\"[agent]:\", event.content.parts[0].text)\n",
        "\n",
        "        print(\"Agent run finished. Exiting context stack...\")\n",
        "        # Main stack cleanup happens automatically here\n",
        "\n",
        "    print(\"Context stack closed, connections cleaned up.\")\n",
        "    return \"Run completed\"  # Or other appropriate return value"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "93f0e0242ea2"
      },
      "outputs": [],
      "source": [
        "# Create server parameters for stdio connection\n",
        "weather_server_params = StdioServerParameters(\n",
        "    command=\"python\",\n",
        "    # Make sure to update to the full absolute path to your weather_server.py file\n",
        "    args=[\"./server/weather_server.py\"],\n",
        ")\n",
        "\n",
        "ct_server_params = StdioServerParameters(\n",
        "    command=\"python\",\n",
        "    args=[\"./server/cocktail2.py\"],\n",
        ")\n",
        "bnb_server_params = StdioServerParameters(\n",
        "    command=\"npx\", args=[\"-y\", \"@openbnb/mcp-server-airbnb\", \"--ignore-robots-txt\"]\n",
        ")\n",
        "\n",
        "server_config_dict = {\n",
        "    \"weather\": weather_server_params,\n",
        "    \"bnb\": bnb_server_params,\n",
        "    \"ct\": ct_server_params,\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b6d51a70a821"
      },
      "outputs": [],
      "source": [
        "await run_multi_agent_with_mcp_clients(\n",
        "    server_config_dict,\n",
        "    \"I want to book an airbnb apartment in LA, CA for 2 nights. 04/28 - 04/30, 2025, two adults, no kid\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "df84795f31cd"
      },
      "source": [
        "### References:\n",
        "https://google.github.io/adk-docs/\n",
        "https://modelcontextprotocol.io/introduction  \n",
        "https://github.com/modelcontextprotocol/python-sdk"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "develop_mcp_with_gemini_and_adk.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
