{
 "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": [
    "# Building an ADK agent using QWEN 3 on Vertex AI\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/open-models/agents/qwen3_adk_vertexai.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%2Fopen-models%2Fagents%2Fqwen3_adk_vertexai.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/open-models/agents/qwen3_adk_vertexai.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/open-models/agents/qwen3_adk_vertexai.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/open-models/agents/qwen3_adk_vertexai.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/open-models/agents/qwen3_adk_vertexai.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/open-models/agents/qwen3_adk_vertexai.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/open-models/agents/qwen3_adk_vertexai.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/open-models/agents/qwen3_adk_vertexai.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",
    "| [Firstname Lastname](https://github.com/username) |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tvgnzT1CKxrO"
   },
   "source": [
    "## Overview\n",
    "\n",
    "This tutorial demonstrates how to build an AI agent using Agent Development Kit (ADK) with Qwen3 model deployed on Vertex AI. Also, it shows how to deploy the agent as a scalable, production-ready service with Vertex AI Agent Engine.\n",
    "\n",
    "What you'll learn how to\n",
    "\n",
    "\n",
    "*   Deploy an open-source model like Qwen3 from Vertex AI Model Garden to a dedicated endpoint.\n",
    "\n",
    "*   Build a tool-calling agent using Google's Agent Development Kit (ADK).\n",
    "\n",
    "*   Integrate a custom model endpoint with ADK using the LiteLLM library.\n",
    "\n",
    "*   Define tools and provide instructions for the agent to use them correctly.\n",
    "\n",
    "*   Test the agent's functionality in a local environment.\n",
    "\n",
    "*   Package and deploy the ADK agent to Vertex AI Agent Engine for a scalable, production-ready implementation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "61RBz8LLbxCR"
   },
   "source": [
    "## Get started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "No17Cw5hgx12"
   },
   "source": [
    "### Install required packages\n",
    "\n",
    "First things first, let's get our toolkit ready. We'll need to install the necessary Python packages for interacting with Vertex AI, ADK, and LiteLLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tFy3H3aPgx12"
   },
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet \"google-cloud-aiplatform[adk,agent_engines]>=1.101.0\" \"litellm>=1.73.6.post1\" \"nest-asyncio>=1.6.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).\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",
    "import os\n",
    "import uuid\n",
    "\n",
    "import vertexai\n",
    "\n",
    "# Project settings\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",
    "if not PROJECT_ID:\n",
    "    raise ValueError(\"PROJECT_ID not set. Please set it in the cell above.\")\n",
    "\n",
    "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
    "BUCKET_URI = \"[your-bucket-uri]\"  # @param {type: \"string\", placeholder: \"[your-bucket-uri]\", isTemplate: true}\n",
    "if not BUCKET_URI or BUCKET_URI == \"[your-bucket-uri]\":\n",
    "    raise ValueError(\"Please set a valid GCS bucket URI in the BUCKET_URI variable.\")\n",
    "\n",
    "# Set config for adk\n",
    "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT_ID\n",
    "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = LOCATION\n",
    "os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"1\"\n",
    "\n",
    "# Agent configuration\n",
    "APP_NAME = \"weather_agent\"\n",
    "USER_ID = \"user_\" + str(uuid.uuid4())\n",
    "SESSION_ID = \"session_\" + str(uuid.uuid4())\n",
    "MAX_TOKENS = 1000\n",
    "TEMPERATURE = 1.0\n",
    "\n",
    "# Initialize Vertex AI\n",
    "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5303c05f7aa6"
   },
   "source": [
    "### Import libraries\n",
    "\n",
    "Now, let's import all the Python libraries and modules we'll be using throughout this tutorial. This includes components from Vertex AI, ADK, LiteLLM, and standard Python libraries.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6fc324893334"
   },
   "outputs": [],
   "source": [
    "import functools\n",
    "import inspect\n",
    "from typing import Callable, Optional, get_type_hints\n",
    "\n",
    "import google.auth\n",
    "import nest_asyncio\n",
    "from google import adk\n",
    "from google.adk.agents import Agent\n",
    "from google.adk.models.lite_llm import LiteLlm\n",
    "from google.adk.sessions import InMemorySessionService\n",
    "from google.adk.tools import FunctionTool\n",
    "from google.cloud import aiplatform\n",
    "from google.genai import types\n",
    "from litellm import completion\n",
    "from vertexai import agent_engines\n",
    "from vertexai.preview import model_garden, reasoning_engines\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yUr0RADhgshQ"
   },
   "source": [
    "### Helpers\n",
    "\n",
    "We'll define a few helper functions to simplify running our agent and interacting with it in a chat-like interface. These helpers will allow us to test the agent both in the local runner and on the remote Agent Engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Izt5PN6RguyF"
   },
   "outputs": [],
   "source": [
    "def python_type_to_json_schema_type(py_type):\n",
    "    \"\"\"Maps Python types to JSON schema types.\"\"\"\n",
    "    if py_type in (int, float):\n",
    "        return \"integer\" if py_type is int else \"number\"\n",
    "    elif py_type is str:\n",
    "        return \"string\"\n",
    "    elif py_type is bool:\n",
    "        return \"boolean\"\n",
    "    elif py_type is list:\n",
    "        return \"array\"\n",
    "    elif py_type is dict:\n",
    "        return \"object\"\n",
    "    return \"string\"  # Default to string for unknown types\n",
    "\n",
    "\n",
    "def convert_to_openai_tool(func):\n",
    "    \"\"\"\n",
    "    Converts a Python function into an OpenAI tool schema.\n",
    "    \"\"\"\n",
    "    # Get the function's signature and docstring\n",
    "    signature = inspect.signature(func)\n",
    "    docstring = inspect.getdoc(func) or \"\"\n",
    "\n",
    "    # Get type hints\n",
    "    type_hints = get_type_hints(func)\n",
    "\n",
    "    # Initialize the schema\n",
    "    function_schema = {\n",
    "        \"name\": func.__name__,\n",
    "        \"description\": docstring.split(\"\\n\\n\")[0],\n",
    "        \"parameters\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {},\n",
    "            \"required\": [],\n",
    "        },\n",
    "    }\n",
    "\n",
    "    # Process each parameter\n",
    "    for param_name, param in signature.parameters.items():\n",
    "        param_type = type_hints.get(param_name, str)\n",
    "        json_type = python_type_to_json_schema_type(param_type)\n",
    "\n",
    "        function_schema[\"parameters\"][\"properties\"][param_name] = {\n",
    "            \"type\": json_type,\n",
    "            \"description\": f\"The {param_name} for the function.\",\n",
    "        }\n",
    "\n",
    "        # Check if the parameter is required\n",
    "        if param.default is inspect.Parameter.empty:\n",
    "            function_schema[\"parameters\"][\"required\"].append(param_name)\n",
    "\n",
    "    # Final tool schema structure\n",
    "    tool_schema = {\"type\": \"function\", \"function\": function_schema}\n",
    "\n",
    "    return tool_schema\n",
    "\n",
    "\n",
    "def get_base_url(endpoint: aiplatform.Endpoint) -> str:\n",
    "    \"\"\"Get the base URL for the endpoint.\"\"\"\n",
    "    location = endpoint.resource_name.split(\"/\")[3]\n",
    "    if endpoint.gca_resource.dedicated_endpoint_enabled:\n",
    "        return (\n",
    "            f\"https://{endpoint.gca_resource.dedicated_endpoint_dns}\"\n",
    "            f\"/v1beta1/{endpoint.resource_name}\"\n",
    "        )\n",
    "    return (\n",
    "        f\"https://{location}-aiplatform.googleapis.com\"\n",
    "        f\"/v1beta1/{endpoint.resource_name}\"\n",
    "    )\n",
    "\n",
    "\n",
    "def run_single_turn(\n",
    "    runner: adk.Runner, query: str, session_id: str, user_id: str\n",
    ") -> Optional[str]:\n",
    "    \"\"\"Runs one turn with a local ADK runner.\"\"\"\n",
    "    content = types.Content(role=\"user\", parts=[types.Part(text=query)])\n",
    "    events = runner.run(user_id=user_id, session_id=session_id, new_message=content)\n",
    "    for event in events:\n",
    "        if event.is_final_response():\n",
    "            return event.content.parts[0].text\n",
    "    return None\n",
    "\n",
    "\n",
    "def run_remote_single_turn(\n",
    "    remote_app: reasoning_engines.templates.adk.AdkApp,\n",
    "    query: str,\n",
    "    session_id: str,\n",
    "    user_id: str,\n",
    ") -> Optional[str]:\n",
    "    \"\"\"Runs one turn on a deployed Agent Engine.\"\"\"\n",
    "    events = remote_app.stream_query(\n",
    "        user_id=user_id, session_id=session_id, message=query\n",
    "    )\n",
    "    final_event = [event for event in events][-1]\n",
    "    return final_event[\"content\"][\"parts\"][0][\"text\"]\n",
    "\n",
    "\n",
    "def chat_loop(\n",
    "    turn_handler: Callable[[str, str, str], Optional[str]],\n",
    "    session_id: str,\n",
    "    user_id: str,\n",
    ") -> None:\n",
    "    \"\"\"Main chat interface loop. It is given the function to handle each turn.\"\"\"\n",
    "    print(\"\\nStarting chat. Type 'exit' or 'quit' to end.\")\n",
    "    while True:\n",
    "        try:\n",
    "            user_input = input(\"You: \")\n",
    "            if user_input.lower() in [\"quit\", \"exit\", \"bye\"]:\n",
    "                print(\"\\nAssistant: Thank you for chatting. Have a great day!\")\n",
    "                break\n",
    "\n",
    "            response = turn_handler(user_input, session_id, user_id)\n",
    "            if response:\n",
    "                print(f\"Assistant: {response}\")\n",
    "            else:\n",
    "                print(\"Assistant: I received no response.\")\n",
    "        except (KeyboardInterrupt, EOFError):\n",
    "            print(\"\\nAssistant: Chat ended.\")\n",
    "            break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QAAOd02fV43F"
   },
   "source": [
    "## Deploy Qwen3 model using Vertex AI Model Garden SDK\n",
    "\n",
    "Deploy the Qwen3 model from the Vertex AI Model Garden. Model Garden provides a curated collection of foundation models that can be easily deployed to a Vertex AI Endpoint. This gives us a scalable, managed API for our model, which our agent will use as its \"brain.\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l764864DiyCR"
   },
   "outputs": [],
   "source": [
    "# Load the model from Model Garden\n",
    "model = model_garden.OpenModel(\"publishers/qwen/models/qwen3@qwen3-1.7b\")\n",
    "\n",
    "# Get container specification\n",
    "serving_container_spec = model.list_deploy_options()[0].container_spec\n",
    "\n",
    "# Deploy the model\n",
    "endpoint = model.deploy(\n",
    "    serving_container_spec=serving_container_spec,\n",
    "    machine_type=\"g2-standard-12\",\n",
    "    accelerator_type=\"NVIDIA_L4\",\n",
    "    accelerator_count=1,\n",
    "    use_dedicated_endpoint=False,\n",
    "    spot=False,\n",
    "    deploy_request_timeout=1800\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eDiKHwP4bAUh"
   },
   "source": [
    "## Test the deployed endpoint with a simple request using LiteLlm\n",
    "\n",
    "Before building the full agent, it's a good practice to ensure we can communicate with our newly deployed Qwen3 endpoint. We'll use the LiteLLM library for this.\n",
    "\n",
    "LiteLLM provides a simple, unified interface to call over 100 different LLM APIs. It allows us to interact with our custom Vertex AI endpoint using the same syntax as the OpenAI API. This is incredibly useful because the ADK has a built-in LiteLlm model wrapper, making integration seamless.\n",
    "\n",
    "Get the necessary authentication credentials for our Google Cloud project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b0aRsdUma_uv"
   },
   "outputs": [],
   "source": [
    "# Get credentials\n",
    "creds, _ = google.auth.default()\n",
    "auth_req = google.auth.transport.requests.Request()\n",
    "creds.refresh(auth_req)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BtOuTPiwJ4Ow"
   },
   "source": [
    "Define a simple Python function (get_weather) that we'll later turn into a tool for our agent.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DQVqKlVreqSc"
   },
   "outputs": [],
   "source": [
    "# Define a simple tool\n",
    "def get_weather(city: str) -> str:\n",
    "\n",
    "    \"\"\"\n",
    "    Get weather information for a city.\n",
    "\n",
    "    Args:\n",
    "        city: The city name to get weather for.\n",
    "\n",
    "    Returns:\n",
    "        Weather information string.\n",
    "    \"\"\"\n",
    "    city_lower = city.lower()\n",
    "\n",
    "    # Simple weather simulation\n",
    "    weather_data = {\n",
    "        \"san francisco\": \"It's 70°F (21°C) and foggy.\",\n",
    "        \"sf\": \"It's 70°F (21°C) and foggy.\",\n",
    "        \"new york\": \"It's 75°F (24°C) and partly cloudy.\",\n",
    "        \"nyc\": \"It's 75°F (24°C) and partly cloudy.\",\n",
    "        \"london\": \"It's 60°F (15°C) and rainy.\",\n",
    "        \"tokyo\": \"It's 72°F (22°C) and sunny.\",\n",
    "    }\n",
    "\n",
    "    for key, weather in weather_data.items():\n",
    "        if key in city_lower:\n",
    "            return weather\n",
    "\n",
    "    return f\"It's 80°F (27°C) and sunny in {city}.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-QNb9NnlJ7Mh"
   },
   "source": [
    "Use litellm.completion to send a direct request to our Qwen3 endpoint (vertex_ai/openai/{endpoint.name}). This verifies that the endpoint is live and that our connection is configured correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KeuJAfKmdBuy"
   },
   "outputs": [],
   "source": [
    "response = completion(\n",
    "    model=f\"vertex_ai/openai/{endpoint.name}\",\n",
    "    messages=[{\"content\": \"Hello, how are you?/no_think\", \"role\": \"user\"}],\n",
    "    tools=[\n",
    "        # Convert the function to an OpenAI tool schema\n",
    "        convert_to_openai_tool(get_weather)\n",
    "    ],\n",
    "    temperature=0.7,\n",
    "    max_tokens=8192,\n",
    "    extra_body={\n",
    "        \"chat_template_kwargs\": {\"enable_thinking\": True},\n",
    "    },\n",
    ")\n",
    "\n",
    "print(response[\"choices\"][0][\"message\"][\"content\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EdvJRUWRNGHE"
   },
   "source": [
    "## Create an Agent with ADK\n",
    "\n",
    "Now that we have a working model endpoint, it's time to build our agent using the Agent Development Kit. An ADK Agent is defined by several key components:\n",
    "\n",
    "\n",
    "*  model: The LLM the agent will use. Here, we'll use ADK's LiteLlm wrapper to point to our Qwen3 endpoint.\n",
    "*   name: A unique identifier for the agent.\n",
    "\n",
    "*   description: A high-level summary of what the agent does.\n",
    "*   instruction: The core prompt that guides the agent's behavior. This is where we tell the agent how and when to use its tools.\n",
    "*   tools: A list of functions the agent can call to interact with external systems. We'll use the FunctionTool wrapper to make our get_weather function available to the agent.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7QHxMswxfTLW"
   },
   "outputs": [],
   "source": [
    "root_agent = Agent(\n",
    "    model=LiteLlm(\n",
    "        model=f\"vertex_ai/openai/{endpoint.name}\",\n",
    "    ),\n",
    "    name=\"weather_agent\",\n",
    "    description=\"Provides weather information for specific cities.\",\n",
    "    instruction=\"You are a helpful weather assistant. \"\n",
    "    \"When the user asks for the weather in a specific city, \"\n",
    "    \"use the 'get_weather' tool to find the information. \"\n",
    "    \"If the tool returns an error, inform the user politely. \"\n",
    "    \"If the tool is successful, present the weather report clearly.\"\n",
    "    \"After using tools, provide a natural language response incorporating the results.\"\n",
    "    \"/no_think\",\n",
    "    tools=[FunctionTool(func=get_weather)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VScXms74fxV5"
   },
   "source": [
    "## Test the agent locally\n",
    "\n",
    "Before deploying our agent to a production environment, it's crucial to test its logic and behavior locally. The ADK Runner allows us to do this easily.\n",
    "\n",
    "The Runner orchestrates the interaction between the user, the agent, and the various services. For this local test, we'll use an InMemorySessionService to keep track of the conversation history.\n",
    "\n",
    "The chat_loop helper function will provide a simple command-line interface to interact with our agent and see how it responds and uses its get_weather tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZzFNqH8kf07m"
   },
   "outputs": [],
   "source": [
    "# Create a session to store conversations\n",
    "session_service = InMemorySessionService()\n",
    "\n",
    "# Create the specific session where the conversation will happen\n",
    "session = await session_service.create_session(\n",
    "    app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID\n",
    ")\n",
    "\n",
    "# Create the runner to run the agent\n",
    "runner = adk.Runner(\n",
    "    agent=root_agent, app_name=APP_NAME, session_service=session_service\n",
    ")\n",
    "\n",
    "local_turn_handler = functools.partial(run_single_turn, runner)\n",
    "\n",
    "chat_loop(local_turn_handler, session.id, USER_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6XRsZZpqhFES"
   },
   "source": [
    "## Deploy the agent on Agent Engine\n",
    "\n",
    "Once we're satisfied with our agent's local performance, the final step is to deploy it as a scalable, production-ready service using Vertex AI Agent Engine. This powerful platform takes our ADK-defined agent and hosts it as a managed endpoint, handling all the underlying infrastructure, scaling, and logging."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VFIInWqpXJs6"
   },
   "source": [
    "### Prepare your agent for Agent Engine\n",
    "\n",
    "To prepare our agent for deployment, we wrap it in a reasoning_engines.AdkApp. This class serves as the bridge between the ADK framework and the Agent Engine deployment platform.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iJxOA631hH1a"
   },
   "outputs": [],
   "source": [
    "app = reasoning_engines.AdkApp(\n",
    "    agent=root_agent,\n",
    "    enable_tracing=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WuqACpAXXZnJ"
   },
   "source": [
    "### Send queries to your agent (locally)\n",
    "\n",
    "Before pushing to the cloud, we can use the AdkApp object to run a final local test. This ensures that the agent is packaged correctly and behaves as expected within the Agent Engine wrapper.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P74ol0xtXbC-"
   },
   "outputs": [],
   "source": [
    "session = app.create_session(user_id=USER_ID)\n",
    "\n",
    "local_app_turn_handler = functools.partial(run_remote_single_turn, app)\n",
    "\n",
    "chat_loop(local_app_turn_handler, session.id, USER_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eVjLyb-8X9o3"
   },
   "source": [
    "### Deploy your agent to Agent Engine\n",
    "\n",
    "Deploying is as simple as calling the agent_engines.create() function. We pass our AdkApp object and specify any Python package requirements. Agent Engine handles the rest: it containerizes the application, provisions the necessary resources, and exposes a secure endpoint for our agent.\n",
    "\n",
    "Note: The deployment process can take a few minutes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-1khvUTUX_F2"
   },
   "outputs": [],
   "source": [
    "remote_app = agent_engines.create(\n",
    "    agent_engine=app,\n",
    "    requirements=[\n",
    "        \"google-cloud-aiplatform[adk,agent_engines]>=1.101.0\",\n",
    "        \"litellm==1.73.6.post1\",\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S2Hm49QIYhWd"
   },
   "source": [
    "### Send queries to your agent (remote)\n",
    "\n",
    "With the agent deployed, we get a remote_app object that points to our live Agent Engine endpoint. We can interact with this remote agent using the exact same chat_loop function, demonstrating the seamless transition from local development to cloud deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "crfG5FqcYjzy"
   },
   "outputs": [],
   "source": [
    "remote_session = remote_app.create_session(user_id=USER_ID)\n",
    "\n",
    "remote_app_turn_handler = functools.partial(run_remote_single_turn, remote_app)\n",
    "\n",
    "chat_loop(remote_app_turn_handler, remote_session.id, USER_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2a4e033321ad"
   },
   "source": [
    "## Cleaning up\n",
    "\n",
    "To avoid incurring ongoing charges to your Google Cloud account, it's important to clean up the resources you've created. The following code will delete the deployed Agent Engine and the Vertex AI Endpoint for the Qwen3 model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kJpYPexsYyCD"
   },
   "outputs": [],
   "source": [
    "delete_endpoint = False\n",
    "delete_agent_engine = False\n",
    "\n",
    "if delete_endpoint:\n",
    "    endpoint.delete(force=True)\n",
    "\n",
    "if delete_agent_engine:\n",
    "    remote_app.delete(force=True)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "qwen3_adk_vertexai.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "python3",
   "name": "workbench-notebooks.m130",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m130"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
