{
 "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": "JAPoU8Sm5E6e"
   },
   "source": [
    "# Claude with ADK 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_claude_with_adk_on_agent_engine.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_claude_with_adk_on_agent_engine.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_claude_with_adk_on_agent_engine.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_claude_with_adk_on_agent_engine.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/agents/agent_engine/tutorial_claude_with_adk_on_agent_engine.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_claude_with_adk_on_agent_engine.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_claude_with_adk_on_agent_engine.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_claude_with_adk_on_agent_engine.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_claude_with_adk_on_agent_engine.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",
    "| [Ivan Nardini](https://github.com/inardini) |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tvgnzT1CKxrO"
   },
   "source": [
    "## Overview\n",
    "\n",
    "This tutorial demonstrates how to build a production-ready, intelligent agent from the ground up. We will use the Google Agent Development Kit (ADK) to orchestrate a powerful third-party model (Anthropic's Claude), connect it to live data with external tools using Model Context Protocol (MCP), and deploy it as a scalable application on Vertex AI Agent Engine.\n",
    "\n",
    "Finally, we will grant our agent a persistent, long-term memory using Vertex AI Memory Bank and enable tracing capabilities just using Vertex AI Agent Engine.\n",
    "\n",
    "### What you'll learn\n",
    "\n",
    "By the end of this tutorial, you will be able to:\n",
    "*   Build and test a local agent using the ADK and a third-party model like Claude.\n",
    "*   Augment an agent with tools to interact with external APIs using MCP (e.g., Reddit).\n",
    "*   Deploy your agent as a scalable, managed application using Vertex AI Agent Engine.\n",
    "*   Enable tracing and long-term memory with Vertex AI Memory Bank and the `PreloadMemoryTool`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "61RBz8LLbxCR"
   },
   "source": [
    "## Get started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "No17Cw5hgx12"
   },
   "source": [
    "### Install required packages\n",
    "\n",
    "Our project relies on several key libraries.\n",
    "\n",
    "The `google-cloud-aiplatform` SDK provides the tools for Agent Engine and memory, `google-adk` is the core agent framework, and `anthropic` allows us to use Claude models on Vertex AI.\n",
    "\n",
    "The other packages support the Reddit tool we'll add later.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tFy3H3aPgx12"
   },
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet \"google-cloud-aiplatform[adk,agent-engines]>=1.110.0\" \"anthropic[vertex]>=0.55.0\" \"google-adk[eval]>=1.4.2\" \"litellm>=1.73.1\" \"fastmcp>=0.2.0\" \"redditwarp>=1.0.0\" \"dnspython>=2.7.0\" \"praw>=7.8.1\" \"vertexai>=1.43.0\" \"click>=8.2.1\" \"httpx>=0.28.1\" \"uvicorn>=0.34.3\" \"nest-asyncio>=1.6.0\" \"google-auth>=2.40.3\" \"aiofiles>=24.1.0\""
   ]
  },
  {
   "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": "DF4l8DTdWgPY"
   },
   "source": [
    "### Set Google Cloud project information\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). Also you can learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).\n",
    "\n",
    "Let's set our project, location, and a unique Cloud Storage bucket name. This bucket will be used by Vertex AI for staging artifacts during deployment."
   ]
  },
  {
   "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",
    "import os\n",
    "import vertexai\n",
    "\n",
    "GOOGLE_CLOUD_PROJECT = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
    "\n",
    "if not GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT == \"[your-project-id]\":\n",
    "    GOOGLE_CLOUD_PROJECT = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
    "\n",
    "GOOGLE_CLOUD_LOCATION = \"europe-west1\"\n",
    "\n",
    "GOOGLE_CLOUD_BUCKET_NAME = \"[your-bucket-name]\"  # @param {type: \"string\", placeholder: \"[your-bucket-name]\", isTemplate: true}\n",
    "GOOGLE_CLOUD_BUCKET = f\"gs://{GOOGLE_CLOUD_BUCKET_NAME}\"\n",
    "\n",
    "# Create the bucket if it doesn't exist\n",
    "! gsutil mb -l $GOOGLE_CLOUD_LOCATION -p $GOOGLE_CLOUD_PROJECT $GOOGLE_CLOUD_BUCKET\n",
    "\n",
    "# Initialize Vertex AI client\n",
    "vertexai.init(\n",
    "    project=GOOGLE_CLOUD_PROJECT,\n",
    "    location=GOOGLE_CLOUD_LOCATION,\n",
    "    staging_bucket=GOOGLE_CLOUD_BUCKET,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5303c05f7aa6"
   },
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6fc324893334"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import logging\n",
    "logging.getLogger(\"google_adk\").setLevel(logging.CRITICAL)\n",
    "logging.getLogger(\"asyncio\").setLevel(logging.CRITICAL)\n",
    "import uuid\n",
    "from typing import Any, Iterator, Optional\n",
    "import aiofiles\n",
    "from google.adk.agents import LlmAgent\n",
    "from google.adk.models.anthropic_llm import Claude\n",
    "from google.adk.models.registry import LLMRegistry\n",
    "from google.adk.runners import Runner\n",
    "from google.adk.sessions import InMemorySessionService\n",
    "from google.adk.tools.mcp_tool import StdioConnectionParams\n",
    "from google.adk.tools.mcp_tool.mcp_toolset import (MCPToolset,\n",
    "                                                   StdioServerParameters)\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e43229f3ad4f"
   },
   "source": [
    "### Set variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cf93d5f0ce00"
   },
   "outputs": [],
   "source": [
    "# Model Selection\n",
    "MODEL_ID = \"claude-sonnet-4@20250514\"\n",
    "\n",
    "# (Optional) Set Reddit Credentials\n",
    "REDDIT_CLIENT_ID = \"[your-reddit-client-id]\"  # @param {type: \"string\", placeholder: \"[your-reddit-client-id]\", isTemplate: true}\n",
    "REDDIT_CLIENT_SECRET = \"[your-reddit-client-secret]\"  # @param {type: \"string\", placeholder: \"[your-reddit-client-secret]\", isTemplate: true}\n",
    "REDDIT_REFRESH_TOKEN = \"[your-reddit-refresh-token]\"  # @param {type: \"string\", placeholder: \"[your-reddit-refresh-token]\", isTemplate: true}\n",
    "\n",
    "os.environ[\"MODEL_ID\"] = MODEL_ID\n",
    "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = GOOGLE_CLOUD_PROJECT\n",
    "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = GOOGLE_CLOUD_LOCATION\n",
    "\n",
    "if not REDDIT_CLIENT_ID or REDDIT_CLIENT_ID == \"[your-reddit-client-id]\" or not REDDIT_CLIENT_SECRET or REDDIT_CLIENT_SECRET == \"[your-reddit-client-secret]\" or not REDDIT_REFRESH_TOKEN or REDDIT_REFRESH_TOKEN == \"[your-reddit-refresh-token]\":\n",
    "    REDDIT_CLIENT_ID = str(os.environ.get(\"REDDIT_CLIENT_ID\"))\n",
    "    REDDIT_CLIENT_SECRET = str(os.environ.get(\"REDDIT_CLIENT_SECRET\"))\n",
    "    REDDIT_REFRESH_TOKEN = str(os.environ.get(\"REDDIT_REFRESH_TOKEN\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UEJr-QF0NIsj"
   },
   "source": [
    "### Helpers\n",
    "\n",
    "This tutorial uses a helper function, `chat_loop`, to provide a consistent interactive chat interface for testing the agent in both local and deployed environments. This function abstracts the differences between interacting with a local ADK Runner and a remote AgentEngine instance.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5OKhUttINJ2w"
   },
   "outputs": [],
   "source": [
    "def chat_loop(\n",
    "    app, user_id: Optional[str] = None, session_id: Optional[str] = None\n",
    ") -> None:\n",
    "    \"\"\"Interactive chat loop for AI applications.\"\"\"\n",
    "\n",
    "    # Simple setup\n",
    "    user_id = user_id or f\"u_{uuid.uuid4().hex[:8]}\"\n",
    "\n",
    "    # Handle session based on app type\n",
    "    if isinstance(app, (AdkApp, AgentEngine)):\n",
    "        # Only create session if session_id is not provided\n",
    "        if not session_id:\n",
    "            session = app.create_session(user_id=user_id)\n",
    "            # Handle both dict and object session responses\n",
    "            if isinstance(session, dict):\n",
    "                session_id = session[\"id\"]\n",
    "            else:\n",
    "                session_id = 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",
    "        session_id = session_id or f\"s_{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(\n",
    "            f\"Unsupported app type: {type(app)}. Expected AdkApp, AgentEngine, or Runner.\"\n",
    "        )\n",
    "\n",
    "    _print_startup_info(user_id, session_id)\n",
    "\n",
    "    # Main 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",
    "            response = _get_response_text(query_fn(user_input))\n",
    "            print(response or \"(No response generated)\")\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",
    "            continue\n",
    "\n",
    "    print(\"\\n👋 Goodbye!\")\n",
    "\n",
    "\n",
    "def _print_startup_info(user_id: str, session_id: str) -> None:\n",
    "    \"\"\"Print startup information.\"\"\"\n",
    "    print(\"\\n🚀 Starting chat...\")\n",
    "    print(f\"👤 User ID: {user_id}\")\n",
    "    print(f\"📁 Session ID: {session_id}\")\n",
    "    print(\"💬 Type 'exit' or 'quit' to end.\")\n",
    "    print(\"-\" * 50)\n",
    "\n",
    "\n",
    "def _get_response_text(events: Iterator[Any]) -> str:\n",
    "    \"\"\"Extract response text from event stream.\"\"\"\n",
    "    responses = []\n",
    "\n",
    "    for event in events:\n",
    "        # Handle dict-like events (AgentEngine format)\n",
    "        if isinstance(event, dict):\n",
    "            text = _extract_from_dict_event(event)\n",
    "        # Handle object-like events\n",
    "        else:\n",
    "            text = _extract_from_object_event(event)\n",
    "\n",
    "        if text:\n",
    "            responses.append(text)\n",
    "            # Print streaming text in real-time\n",
    "            print(text, end=\"\", flush=True)\n",
    "\n",
    "    return \"\".join(responses)\n",
    "\n",
    "\n",
    "def _extract_from_dict_event(event: dict) -> Optional[str]:\n",
    "    \"\"\"Extract text from dictionary-style events (AgentEngine format).\"\"\"\n",
    "    # Handle AgentEngine response format\n",
    "    if \"parts\" in event and \"role\" in event:\n",
    "        # Only extract text from model responses, skip function calls/responses\n",
    "        if event.get(\"role\") == \"model\":\n",
    "            parts = event.get(\"parts\", [])\n",
    "            text_parts = []\n",
    "\n",
    "            for part in parts:\n",
    "                # Extract text content, skip function calls\n",
    "                if isinstance(part, dict) and \"text\" in part:\n",
    "                    text_parts.append(part[\"text\"])\n",
    "\n",
    "            return \"\".join(text_parts) if text_parts else None\n",
    "        return None\n",
    "\n",
    "    # Handle other dict formats\n",
    "    content = event.get(\"content\", {})\n",
    "    if isinstance(content, str):\n",
    "        return content\n",
    "\n",
    "    parts = content.get(\"parts\", [])\n",
    "    if not parts:\n",
    "        return None\n",
    "\n",
    "    text_parts = []\n",
    "    for part in parts:\n",
    "        if isinstance(part, dict) and \"text\" in part:\n",
    "            text_parts.append(part[\"text\"])\n",
    "\n",
    "    return \"\".join(text_parts) if text_parts else None\n",
    "\n",
    "\n",
    "def _extract_from_object_event(event: Any) -> Optional[str]:\n",
    "    \"\"\"Extract text from object-style events.\"\"\"\n",
    "    # Handle string content directly\n",
    "    content = getattr(event, \"content\", None)\n",
    "    if isinstance(content, str):\n",
    "        return content\n",
    "\n",
    "    # Handle content with parts\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",
    "    # Handle direct text attribute\n",
    "    return getattr(event, \"text\", None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q7VSoCMrMvnu"
   },
   "source": [
    "## Step 1: Create a simple ADK Social Media Agent\n",
    "\n",
    "Let's start with the simplest possible agent: a creative assistant that can brainstorm social media content."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2H_U59SonVeE"
   },
   "source": [
    "### Define your 1st ADK agent\n",
    "\n",
    "The core of our application is the `LlmAgent`. Think of it as the agent's brain. We define its personality and purpose through its `instruction` (also known as a system prompt) and specify which LLM it should use.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oDFJHhf-MxJa"
   },
   "outputs": [],
   "source": [
    "# Register Claude Model for ADK\n",
    "LLMRegistry.register(Claude)\n",
    "\n",
    "# Define the Social Media Assistant Agent\n",
    "root_agent = LlmAgent(\n",
    "    name=\"SocialMediaAssistant\",\n",
    "    model=os.getenv(\"MODEL_ID\", \"claude-sonnet-4@20250514\"),\n",
    "    description=\"Helps create engaging social media content and strategies.\",\n",
    "    instruction=(\n",
    "        \"You are a creative and knowledgeable Social Media Assistant.\\n\"\n",
    "        \"Your goal is to help users create engaging social media content and develop strategies.\\n\"\n",
    "        \"1. Help users brainstorm content ideas for various social media platforms (Instagram, Twitter, LinkedIn, etc.).\\n\"\n",
    "        \"2. Provide suggestions for captions, hashtags, and post timing.\\n\"\n",
    "        \"3. Offer advice on content strategy, audience engagement, and brand voice.\\n\"\n",
    "        \"4. Help with content planning and scheduling recommendations.\\n\"\n",
    "        \"5. Suggest trending topics and content formats that might work well.\\n\"\n",
    "        \"6. Provide best practices for different social media platforms.\\n\"\n",
    "        \"7. Maintain a professional yet creative tone throughout the conversation.\\n\"\n",
    "        \"8. Ask clarifying questions about the user's target audience, platform preferences, and content goals when needed.\"\n",
    "    ),\n",
    "    tools=[],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AYVLYgLZOkly"
   },
   "source": [
    "### Test the simple agent locally\n",
    "\n",
    "With our agent defined, let's chat with it. For local development, the ADK provides in-memory services (`InMemorySessionService`) that simulate a production environment without any cloud dependencies. This is perfect for rapid iteration and testing.\n",
    "\n",
    "The `Runner` is the engine that orchestrates the conversation, managing the flow between the user, the agent, and any services. We'll use a simple `chat_loop` helper to make the interaction feel natural."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZNXgEjB1OnDH"
   },
   "outputs": [],
   "source": [
    "async def test_agent_locally():\n",
    "\n",
    "    # Create the session service\n",
    "    session_service = InMemorySessionService()\n",
    "\n",
    "    # Create the specific session where the conversation will happen\n",
    "    user_id = \"user_01\"\n",
    "    session = await session_service.create_session(\n",
    "        app_name=\"MyApp\", user_id=user_id\n",
    "    )\n",
    "\n",
    "    # Initialize the Runner with the correct session service\n",
    "    runner = Runner(\n",
    "        agent=root_agent,\n",
    "        app_name=\"MyApp\",\n",
    "        session_service=session_service,\n",
    "    )\n",
    "\n",
    "    try:\n",
    "        chat_loop(runner, user_id, session.id)\n",
    "    except KeyboardInterrupt:\n",
    "        print(\"\\n\\n🛑 Chat interrupted.\")\n",
    "    except Exception as e:\n",
    "        print(f\"\\n❌ Error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jF1Tj74pl7UR"
   },
   "source": [
    "You now have a running, simple, agent powered by Claude. Try asking it for some social media post ideas!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NUHH_fX0Orah"
   },
   "outputs": [],
   "source": [
    "await test_agent_locally()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Inru9DT_P4Hi"
   },
   "source": [
    "## Step 2: Augmenting the agent with tools using MCP\n",
    "\n",
    "Our agent is creative, but it's isolated. To be truly useful, it needs to interact with the outside world. We'll give it the ability to research topics on Reddit.\n",
    "\n",
    "Tools are how agents perform actions. The ADK uses the **Model Context Protocol (MCP)**, a standard for building and connecting tools. We'll use a pre-built MCP tool for Reddit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_XIzMtR_rcUN"
   },
   "source": [
    "### Prepare the tool's runtime environment\n",
    "\n",
    "The Reddit tool is an external Python package. To install it, we can use a simple shell script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-9KwhvQtr0-f"
   },
   "outputs": [],
   "source": [
    "!mkdir -p installation_scripts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q5f_iAu-k9e5"
   },
   "outputs": [],
   "source": [
    "install_local_mcp_file = \"\"\"\n",
    "#!/bin/bash\n",
    "# Exit immediately if a command exits with a non-zero status.\n",
    "set -e\n",
    "\n",
    "echo \"Installing MCP Reddit Server\"\n",
    "\n",
    "# Install uv (a fast Python package manager)\n",
    "apt-get update\n",
    "apt-get install -y curl\n",
    "echo \"Installing uv...\"\n",
    "curl -LsSf https://astral.sh/uv/install.sh | sh\n",
    "\n",
    "# Add uv to PATH for current session\n",
    "export PATH=\"$HOME/.local/bin:$PATH\"\n",
    "\n",
    "# Install the mcp-reddit tool using the command from its documentation\n",
    "echo \"Installing mcp-reddit using uv...\"\n",
    "uv pip install \"git+https://github.com/adhikasp/mcp-reddit.git\" --system\n",
    "echo \"MCP Reddit Server installation complete.\"\n",
    "\"\"\"\n",
    "\n",
    "with open(\"installation_scripts/install_local_mcp.sh\", \"w\") as f:\n",
    "    f.write(install_local_mcp_file)\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OYmyvnTO2j9f"
   },
   "outputs": [],
   "source": [
    "!chmod +x installation_scripts/install_local_mcp.sh && ./installation_scripts/install_local_mcp.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jgTU1HTasbvO"
   },
   "source": [
    "### Update the agent definition\n",
    "\n",
    "Now, we update our `LlmAgent` to include the `MCPToolset`. We also refine its instructions, explicitly telling it about its new capabilities. This is critical—the LLM needs to know what tools it has and when to use them.\n",
    "\n",
    "> **Notice**: To solve the `fileno` error in Colab, we must redirect the tool's error stream to a file that supports asynchronous operations. We will use the aiofiles library for this. This requires our test to run inside an async function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1Lha44dp9CCk"
   },
   "outputs": [],
   "source": [
    "def create_agent(errlog):\n",
    "    root_agent = LlmAgent(\n",
    "        name=\"SocialMediaAssistantWithReddit\",\n",
    "        model=os.getenv(\"MODEL_ID\", \"claude-sonnet-4@20250514\"),\n",
    "        description=\"Helps create engaging social media content with Reddit research capabilities.\",\n",
    "        instruction=(\n",
    "        \"You are a creative Social Media Assistant with Reddit research capabilities.\\n\"\n",
    "        \"Your goal is to help users create engaging social media posts for LinkedIn and X.\\n\\n\"\n",
    "        \"You have access to Reddit tools that can help you:\\n\"\n",
    "        \"1. **fetch_hot_threads**: Get hot/trending threads from any subreddit\\n\"\n",
    "        \"2. **fetch_post_content**: Get detailed post content including comments\\n\\n\"\n",
    "        \"Use these tools to gather the latest trends and create data-driven social media content.\"\n",
    "        ),\n",
    "        tools=[\n",
    "            MCPToolset(\n",
    "                connection_params=StdioConnectionParams(\n",
    "                    server_params=StdioServerParameters(\n",
    "                        command=\"mcp-reddit\",\n",
    "                    ),\n",
    "                ),\n",
    "                errlog=errlog,  # Required only in colab environment\n",
    "            )\n",
    "        ],\n",
    "    )\n",
    "    return root_agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1t6ceQ0IuVPg"
   },
   "source": [
    "### Test the agent locally\n",
    "\n",
    "Now, try a prompt like: \"**What are some hot topics in the r/MachineLearning subreddit?**\" You should see the agent activate its tool to fetch live data from Reddit!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z6UiAGbT-jLp"
   },
   "outputs": [],
   "source": [
    "async def test_agent_locally():\n",
    "\n",
    "    # Create the session service\n",
    "    session_service = InMemorySessionService()\n",
    "\n",
    "    # Create the specific session where the conversation will happen\n",
    "    user_id = \"user_01\"\n",
    "    session = await session_service.create_session(\n",
    "        app_name=\"MyApp\", user_id=user_id\n",
    "    )\n",
    "\n",
    "    # Initialize the Runner with the correct session service\n",
    "    errlog = await aiofiles.open(\n",
    "        \"error.log\", \"w+\"\n",
    "    )  # Required only in colab environment.\n",
    "    root_agent = create_agent(errlog)\n",
    "    runner = Runner(\n",
    "        agent=root_agent,\n",
    "        app_name=\"MyApp\",\n",
    "        session_service=session_service\n",
    "    )\n",
    "\n",
    "    try:\n",
    "        chat_loop(runner, user_id, session.id)\n",
    "    finally:\n",
    "        # Ensure the log file is always closed\n",
    "        await errlog.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7d4c5b3b2f21"
   },
   "outputs": [],
   "source": [
    "await test_agent_locally()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ShSFCOuFgb_I"
   },
   "source": [
    "## Step 3: Scaling up to Vertex AI Agent Engine\n",
    "\n",
    "Local testing is fantastic for development speed, but real applications need to be deployed. Vertex AI Agent Engine is a managed service that hosts, scales, and monitors your agent, turning your Python script into a production-ready API.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iEpQB0kb_Cln"
   },
   "source": [
    "### Create the agent module\n",
    "\n",
    "To deploy, we must package our agent's definition into a Python file, which we'll call `root_agent.py`. This file will contain not only the agent logic but also \"builders\" for cloud services like `VertexAiSessionService`, the production equivalent of `InMemorySessionService`.\n",
    "\n",
    "Notice we wrap our agent definition in an `AdkApp`. This is the entry point that Agent Engine will use to serve our agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UwF1E-twlNdr"
   },
   "outputs": [],
   "source": [
    "root_agent_file = f\"\"\"\n",
    "import os\n",
    "from google.adk.models.anthropic_llm import Claude\n",
    "from google.adk.models.registry import LLMRegistry\n",
    "from google.adk.agents import LlmAgent\n",
    "from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, StdioServerParameters\n",
    "from google.adk.tools.mcp_tool import StdioConnectionParams\n",
    "import vertexai\n",
    "from vertexai.preview.reasoning_engines import AdkApp\n",
    "\n",
    "# Set variable\n",
    "GOOGLE_CLOUD_PROJECT=os.getenv('GOOGLE_CLOUD_PROJECT', '{GOOGLE_CLOUD_PROJECT}')\n",
    "GOOGLE_CLOUD_LOCATION=os.getenv('GOOGLE_CLOUD_LOCATION', '{GOOGLE_CLOUD_LOCATION}')\n",
    "MODEL_ID=os.getenv('MODEL_ID', '{MODEL_ID}')\n",
    "REDDIT_CLIENT_ID=os.getenv('REDDIT_CLIENT_ID', '{REDDIT_CLIENT_ID}')\n",
    "REDDIT_CLIENT_SECRET=os.getenv('REDDIT_CLIENT_SECRET', '{REDDIT_CLIENT_SECRET}')\n",
    "REDDIT_REFRESH_TOKEN=os.getenv('REDDIT_REFRESH_TOKEN', '{REDDIT_REFRESH_TOKEN}')\n",
    "\n",
    "# Initialize Vertex AI SDK\n",
    "vertexai.init(project=GOOGLE_CLOUD_PROJECT, location=GOOGLE_CLOUD_LOCATION)\n",
    "\n",
    "# Register Claude Model for ADK\n",
    "LLMRegistry.register(Claude)\n",
    "\n",
    "# Define session builder\n",
    "def session_service_builder():\n",
    "  # This is needed to ensure InitGoogle and AdkApp setup is called first.\n",
    "  from google.adk.sessions import VertexAiSessionService\n",
    "  return VertexAiSessionService(project=GOOGLE_CLOUD_PROJECT, location=GOOGLE_CLOUD_LOCATION)\n",
    "\n",
    "agent_app = AdkApp(\n",
    "    agent=LlmAgent(\n",
    "        name='SocialMediaAssistantWithReddit',\n",
    "        model=os.getenv('MODEL_ID'),\n",
    "        description='Helps create engaging social media content with Reddit research capabilities.',\n",
    "        instruction=(\n",
    "        'You are a creative Social Media Assistant with Reddit research capabilities.\\\\n'\n",
    "        'Your goal is to help users create engaging social media posts for LinkedIn and X.\\\\n\\\\n'\n",
    "        'You have access to Reddit tools that can help you:\\\\n'\n",
    "        '1. **fetch_hot_threads**: Get hot/trending threads from any subreddit\\\\n'\n",
    "        '2. **fetch_post_content**: Get detailed post content including comments\\\\n\\\\n'\n",
    "        'Use these tools to gather the latest trends and create data-driven social media content.'\n",
    "        ),\n",
    "        tools=[\n",
    "            MCPToolset(\n",
    "                connection_params=StdioConnectionParams(\n",
    "                    server_params=StdioServerParameters(\n",
    "                        command='mcp-reddit',\n",
    "                    ),\n",
    "                    timeout=60,\n",
    "                ),\n",
    "            )\n",
    "        ],\n",
    "    ),\n",
    "    session_service_builder=session_service_builder,\n",
    ")\n",
    "\"\"\"\n",
    "\n",
    "with open(\"root_agent.py\", \"w\") as f:\n",
    "    f.write(root_agent_file)\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ebj7YOeSskCW"
   },
   "source": [
    "### Deploy the Agent\n",
    "\n",
    "Now we call `agent_engines.create()`. This command packages our code, installs dependencies, and deploys it to a secure, scalable environment.\n",
    "\n",
    "*   `agent_engine`: We use a `ModuleAgent` to point to our `root_agent.py` file and the `agent_app` object within it.\n",
    "*   `extra_packages`: We include our agent module and the installation script.\n",
    "*   `env_vars`: We securely pass our Reddit API keys to the production environment.\n",
    "*   `build_options`: We instruct the build process to run our custom installation script.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8E6NRMJu_cGt"
   },
   "outputs": [],
   "source": [
    "remote_app = agent_engines.create(\n",
    "    display_name=\"reddit_assistant_agent\",\n",
    "    description=\"A Reddit assistant agent with MCP\",\n",
    "    agent_engine=agent_engines.ModuleAgent(\n",
    "        module_name=\"root_agent\",\n",
    "        agent_name=\"agent_app\",\n",
    "        # You can get the register operations with app = AdkApp(root_agent, session_service_builder), app.setup, app.register_operations()\n",
    "        register_operations={'': ['get_session', 'list_sessions', 'create_session', 'delete_session'],\n",
    "          'async': ['async_get_session',\n",
    "            'async_list_sessions',\n",
    "            'async_create_session',\n",
    "            'async_delete_session',\n",
    "            'async_add_session_to_memory',\n",
    "            'async_search_memory'],\n",
    "          'stream': ['stream_query', 'streaming_agent_run_with_events'],\n",
    "          'async_stream': ['async_stream_query']}\n",
    "    ),\n",
    "    # Python dependencies for the cloud environment.\n",
    "    requirements=[\n",
    "            \"google-cloud-aiplatform[adk,agent-engines]>=1.110.0\",\n",
    "            \"anthropic[vertex]>=0.55.0\",\n",
    "            \"litellm>=1.73.1\",\n",
    "            \"fastmcp>=0.2.0\",\n",
    "            \"redditwarp>=1.0.0\",\n",
    "            \"dnspython>=2.7.0\",\n",
    "            \"praw>=7.8.1\",\n",
    "            \"python-dotenv>=1.0.0\",\n",
    "            \"vertexai>=1.43.0\",\n",
    "        ],\n",
    "    # Files to include in the deployment package.\n",
    "    extra_packages=[\n",
    "        \"root_agent.py\",\n",
    "        \"installation_scripts/install_local_mcp.sh\",\n",
    "    ],\n",
    "    # Environment variables for the deployed agent.\n",
    "    env_vars={\n",
    "        \"MODEL_ID\" : MODEL_ID,\n",
    "        \"REDDIT_CLIENT_ID\": REDDIT_CLIENT_ID,\n",
    "        \"REDDIT_CLIENT_SECRET\": REDDIT_CLIENT_SECRET,\n",
    "        \"REDDIT_REFRESH_TOKEN\": REDDIT_REFRESH_TOKEN,\n",
    "    },\n",
    "    # Custom build steps.\n",
    "    build_options={\n",
    "        \"installation\": [\n",
    "            \"installation_scripts/install_local_mcp.sh\",\n",
    "        ],\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OWP0Vijjs_0T"
   },
   "source": [
    "### Test the agent remotely\n",
    "\n",
    "Once deployed, you can interact with the `remote_app` object just like you did with the local `Runner`. The interface is identical, but the execution is now happening on Google Cloud."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i7psp3vVBbG4"
   },
   "outputs": [],
   "source": [
    "user_id = \"user_01\"\n",
    "session = remote_app.create_session(user_id=user_id)\n",
    "\n",
    "chat_loop(remote_app, user_id=user_id, session_id=session['id'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lufkbySOYuxD"
   },
   "source": [
    "### Generate memories from the previous conversation\n",
    "\n",
    "You have two methods to add memory from previous conversation or session. You can use:\n",
    "\n",
    "*   `generate_memories` method\n",
    "*   `async_add_session_to_memory` method\n",
    "\n",
    "In this case, we use the `generate_memories` method to better control.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ao8XFBeLY1Gz"
   },
   "outputs": [],
   "source": [
    "agent_engine_name = remote_app.resource_name\n",
    "app_name = remote_app.display_name\n",
    "\n",
    "response = agent_engines.generate_memories(\n",
    "  name=agent_engine_name,\n",
    "  vertex_session_source={\n",
    "    \"session\": f\"{agent_engine_name}/sessions/{session['id']}\"\n",
    "  },\n",
    "  scope= {\n",
    "    \"user_id\": user_id,\n",
    "    \"app_name\": app_name\n",
    "  }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q00P1xv2h3Iy"
   },
   "outputs": [],
   "source": [
    "if response.response.generated_memories:\n",
    "    for id, m in enumerate(response.response.generated_memories):\n",
    "        memory = agent_engines.get_memory(name=m.memory.name)\n",
    "        print(\"=========================\")\n",
    "        print(f\"Memory {id}\")\n",
    "        print(\"=========================\")\n",
    "        print(memory.scope)\n",
    "        print(memory.fact)\n",
    "        print(memory)\n",
    "        print(\"\\n\")\n",
    "else:\n",
    "    print(\"No generated memories found.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XtsGh8DpWnZN"
   },
   "source": [
    "## (Optional) Step 4: Enable memory and tracing capabilities with your agent\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HBhCKLcUY1aU"
   },
   "source": [
    "### Update the agent module\n",
    "\n",
    "We'll update our `root_agent.py` to include two new components:\n",
    "\n",
    "1.  `VertexAiMemoryBankService`: The service that connects our agent to the Memory Bank backend.\n",
    "2.  `PreloadMemoryTool`: A built-in ADK tool that automatically queries the Memory Bank at the start of each turn, retrieving relevant facts about the user.\n",
    "\n",
    "We simply add a `memory_service_builder` and include the `PreloadMemoryTool` in our agent's tool list. We'll also set `enable_tracing=True` for better observability in the Google Cloud console."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WBnuuvx_Wmx2"
   },
   "outputs": [],
   "source": [
    "agent_engine_id = remote_app.resource_name.split(\"/\")[-1]\n",
    "\n",
    "root_agent_file = f\"\"\"\n",
    "import os\n",
    "from google import adk\n",
    "from google.adk.models.anthropic_llm import Claude\n",
    "from google.adk.models.registry import LLMRegistry\n",
    "from google.adk.agents import LlmAgent\n",
    "from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, StdioServerParameters\n",
    "from google.adk.tools.mcp_tool import StdioConnectionParams\n",
    "import vertexai\n",
    "from vertexai.preview.reasoning_engines import AdkApp\n",
    "\n",
    "# Set variable\n",
    "GOOGLE_CLOUD_PROJECT=os.getenv('GOOGLE_CLOUD_PROJECT', '{GOOGLE_CLOUD_PROJECT}')\n",
    "GOOGLE_CLOUD_LOCATION=os.getenv('GOOGLE_CLOUD_LOCATION', '{GOOGLE_CLOUD_LOCATION}')\n",
    "GOOGLE_AGENT_ENGINE_ID=os.getenv('GOOGLE_AGENT_ENGINE_ID', '{agent_engine_id}')\n",
    "MODEL_ID=os.getenv('MODEL_ID', '{MODEL_ID}')\n",
    "REDDIT_CLIENT_ID=os.getenv('REDDIT_CLIENT_ID', '{REDDIT_CLIENT_ID}')\n",
    "REDDIT_CLIENT_SECRET=os.getenv('REDDIT_CLIENT_SECRET', '{REDDIT_CLIENT_SECRET}')\n",
    "REDDIT_REFRESH_TOKEN=os.getenv('REDDIT_REFRESH_TOKEN', '{REDDIT_REFRESH_TOKEN}')\n",
    "\n",
    "\n",
    "# Initialize Vertex AI SDK\n",
    "vertexai.init(project=GOOGLE_CLOUD_PROJECT, location=GOOGLE_CLOUD_LOCATION)\n",
    "\n",
    "# Register Claude Model for ADK\n",
    "LLMRegistry.register(Claude)\n",
    "\n",
    "# Define session builder\n",
    "def session_service_builder():\n",
    "  from google.adk.sessions import VertexAiSessionService\n",
    "  return VertexAiSessionService(project=GOOGLE_CLOUD_PROJECT, location=GOOGLE_CLOUD_LOCATION)\n",
    "\n",
    "# Define memory builder\n",
    "def memory_service_builder():\n",
    "    from google.adk.memory import VertexAiMemoryBankService\n",
    "    return VertexAiMemoryBankService(project=GOOGLE_CLOUD_PROJECT, location=GOOGLE_CLOUD_LOCATION, agent_engine_id=GOOGLE_AGENT_ENGINE_ID)\n",
    "\n",
    "\n",
    "agent_app = AdkApp(\n",
    "    agent=LlmAgent(\n",
    "        name='SocialMediaAssistantWithReddit',\n",
    "        model=os.getenv('MODEL_ID'),\n",
    "        description='Helps create engaging social media content with Reddit research capabilities.',\n",
    "        instruction=(\n",
    "        'You are a creative Social Media Assistant with Reddit research capabilities.\\\\n'\n",
    "        'Your goal is to help users create engaging social media posts for LinkedIn and X.\\\\n\\\\n'\n",
    "        'You have access to Reddit tools that can help you:\\\\n'\n",
    "        '1. **fetch_hot_threads**: Get hot/trending threads from any subreddit\\\\n'\n",
    "        '2. **fetch_post_content**: Get detailed post content including comments\\\\n\\\\n'\n",
    "        'Use these tools to gather the latest trends and create data-driven social media content.'\n",
    "        ),\n",
    "        tools=[\n",
    "            MCPToolset(\n",
    "                connection_params=StdioConnectionParams(\n",
    "                    server_params=StdioServerParameters(\n",
    "                        command='mcp-reddit',\n",
    "                    ),\n",
    "                    timeout=60,\n",
    "                ),\n",
    "            ),\n",
    "            adk.tools.preload_memory_tool.PreloadMemoryTool()\n",
    "        ]\n",
    "    ),\n",
    "    session_service_builder=session_service_builder,\n",
    "    memory_service_builder=memory_service_builder,\n",
    "    enable_tracing=True\n",
    ")\n",
    "\"\"\"\n",
    "\n",
    "with open(\"root_agent.py\", \"w\") as f:\n",
    "    f.write(root_agent_file)\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0_ua1MQitTgK"
   },
   "source": [
    "### Update the deployed agent\n",
    "\n",
    "\n",
    "Instead of creating a new agent, we use `agent_engines.update()` to deploy our new version. The process is the same, but it updates the existing endpoint.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VoxE2IErP8UD"
   },
   "outputs": [],
   "source": [
    "remote_app = agent_engines.update(\n",
    "    resource_name=agent_engine_name,\n",
    "    agent_engine=agent_engines.ModuleAgent(\n",
    "        module_name=\"root_agent\",\n",
    "        agent_name=\"agent_app\",\n",
    "        register_operations={'': ['get_session', 'list_sessions', 'create_session', 'delete_session'],\n",
    "          'async': ['async_get_session',\n",
    "            'async_list_sessions',\n",
    "            'async_create_session',\n",
    "            'async_delete_session',\n",
    "            'async_add_session_to_memory',\n",
    "            'async_search_memory'],\n",
    "          'stream': ['stream_query', 'streaming_agent_run_with_events'],\n",
    "          'async_stream': ['async_stream_query']}\n",
    "    ),\n",
    "    requirements=[\n",
    "            \"google-cloud-aiplatform[adk,agent-engines]>=1.110.0\",\n",
    "            \"anthropic[vertex]>=0.55.0\",\n",
    "            \"litellm>=1.73.1\",\n",
    "            \"fastmcp>=0.2.0\",\n",
    "            \"redditwarp>=1.0.0\",\n",
    "            \"dnspython>=2.7.0\",\n",
    "            \"praw>=7.8.1\",\n",
    "            \"python-dotenv>=1.0.0\",\n",
    "            \"vertexai>=1.43.0\",\n",
    "            \"aiofiles>=24.1.0\"\n",
    "        ],\n",
    "    extra_packages=[\n",
    "        \"root_agent.py\",\n",
    "        \"installation_scripts/install_local_mcp.sh\",\n",
    "    ],\n",
    "    env_vars={\n",
    "        \"MODEL_ID\" : MODEL_ID,\n",
    "        \"REDDIT_CLIENT_ID\": REDDIT_CLIENT_ID,\n",
    "        \"REDDIT_CLIENT_SECRET\": REDDIT_CLIENT_SECRET,\n",
    "        \"REDDIT_REFRESH_TOKEN\": REDDIT_REFRESH_TOKEN,\n",
    "        \"GOOGLE_AGENT_ENGINE_ID\" : agent_engine_id,\n",
    "    },\n",
    "    build_options={\n",
    "        \"installation\": [\n",
    "            \"installation_scripts/install_local_mcp.sh\",\n",
    "        ],\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QqKrlm_rtgD5"
   },
   "source": [
    "### Test the agent remotely\n",
    "\n",
    "Once updated, you can interact as before. But in this case the agent should remember details about previous conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "53JDho2ij_m-"
   },
   "outputs": [],
   "source": [
    "user_id = \"user_01\"\n",
    "session = remote_app.create_session(user_id=user_id)\n",
    "\n",
    "chat_loop(remote_app, user_id=user_id, session_id=session['id'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2a4e033321ad"
   },
   "source": [
    "## Cleaning up\n",
    "\n",
    "To avoid incurring future charges, it's important to delete the resources we've created. This command will delete the Agent Engine deployment and the associated services. We also remove the Cloud Storage bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3fB6j7uMmSl7"
   },
   "outputs": [],
   "source": [
    "delete_agent_engine = True\n",
    "delete_bucket = True\n",
    "\n",
    "if delete_agent_engine:\n",
    "    agent_engine_name = remote_app.delete(force=True)\n",
    "\n",
    "if delete_bucket:\n",
    "    !gsutil rm -r $GOOGLE_CLOUD_BUCKET\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "tutorial_claude_with_adk_on_agent_engine.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "python3",
   "name": "workbench-notebooks.m131",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m131"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
