{
  "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": [
        "# Getting Started with Live API on Agent Engine\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/agents/agent_engine/tutorial_get_started_with_live_api_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_get_started_with_live_api_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_get_started_with_live_api_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_get_started_with_live_api_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_get_started_with_live_api_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_get_started_with_live_api_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_get_started_with_live_api_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_get_started_with_live_api_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_get_started_with_live_api_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",
        "| Hanfei Sun, Huang Xia, [Ivan Nardini](https://github.com/inardini) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to build, deploy, and interact with **bidirectional streaming agents** using **Vertex AI Agent Engine** and the **Live API**.\n",
        "\n",
        "Bidirectional (bidi) streaming represents a paradigm shift in AI interactions. Instead of the traditional request-response pattern, it enables **real-time, two-way communication** where both user and AI can exchange data simultaneously.\n",
        "\n",
        "Agent Engine is Google's fully-managed, serverless platform for running AI agents at scale. Combined with the Live API's real-time streaming capabilities, you can create agents that engage in natural, fluid conversations with support for text, audio, and video.\n",
        "\n",
        "In this tutorial, you will:\n",
        "\n",
        "* **Build** three different types of streaming agents with increasing complexity\n",
        "* **Understand** bidirectional streaming patterns and when to use each one\n",
        "* **Deploy** agents to Agent Engine with a single command\n",
        "* **Interact** with agents\n",
        "* **Implement** real-time audio conversations"
      ]
    },
    {
      "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",
        "\n",
        "First, install the necessary packages for building streaming agents with Live API support.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet --force-reinstall git+https://github.com/googleapis/python-aiplatform.git\n",
        "%pip install --upgrade --quiet numpy websockets 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": "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 vertexai\n",
        "\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\", \"global\")\n",
        "\n",
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type: \"string\", placeholder: \"[your-bucket-name]\", isTemplate: true}\n",
        "if not BUCKET_NAME or BUCKET_NAME == \"[your-bucket-name]\":\n",
        "    BUCKET_NAME = PROJECT_ID\n",
        "\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "\n",
        "MODEL_ID = \"gemini-2.0-flash-live-preview-04-09\"\n",
        "\n",
        "# Set enviroment variables 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",
        "# Create the staging bucket for Agent Engine\n",
        "!gsutil mb -p $PROJECT_ID -l $LOCATION $BUCKET_URI\n",
        "\n",
        "# Initialize Vertex AI\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)\n",
        "client = vertexai.Client(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "print(f\"Project ID: {PROJECT_ID}\")\n",
        "print(f\"Location: {LOCATION}\")\n",
        "print(f\"Staging Bucket: {BUCKET_URI}\")\n",
        "print(f\"Model ID: {MODEL_ID}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "Here we import all the necessary Python classes and functions we'll use throughout the tutorial. The imports are organized by category for clarity.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6fc324893334"
      },
      "outputs": [],
      "source": [
        "# Standard library\n",
        "import os\n",
        "import asyncio\n",
        "import json\n",
        "import base64\n",
        "import time\n",
        "from typing import Any, Dict, AsyncIterator, Iterator\n",
        "\n",
        "# Google Cloud and Vertex AI\n",
        "import vertexai\n",
        "from vertexai.preview.reasoning_engines import AdkApp\n",
        "from vertexai import types as vertexai_types\n",
        "import google.auth\n",
        "import google.auth.transport.requests\n",
        "\n",
        "# ADK (Agent Development Kit)\n",
        "from google.adk import LlmAgent\n",
        "from google.adk.tools.tool_context import ToolContext\n",
        "from google.adk.sessions.in_memory_session_service import InMemorySessionService\n",
        "from google.adk.memory.in_memory_memory_service import InMemoryMemoryService\n",
        "from google.adk.agents.live_request_queue import LiveRequest\n",
        "from google.adk.events import Event\n",
        "from google.genai import types\n",
        "\n",
        "# Audio and WebSocket handling\n",
        "import numpy as np\n",
        "import websockets\n",
        "\n",
        "# Notebook display utilities\n",
        "from IPython.display import Audio, display\n",
        "import logging; logging.getLogger().setLevel(logging.INFO)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4aFYrOd0qvMc"
      },
      "source": [
        "## Building your first Echo Agent\n",
        "\n",
        "Before start using the Live API, let's start with a simple echo agent to understand the fundamental patterns of Agent Engine. This agent will demonstrate the three types of interactions supported: standard query, streaming query, and bidirectional streaming.\n",
        "\n",
        "### Understanding Agent Structure\n",
        "\n",
        "Every Agent Engine agent follows a specific pattern with three key methods:\n",
        "\n",
        "1. **Init method**: This must be lightweight and pickle-able (serializable), so we only store simple config values.\n",
        "2. **Setup method**: This is where heavy initialization (like creating database clients or loading models) goes. Agent Engine calls this once when the serverless container spins up.\n",
        "3. **Query methods**: Handle different interaction patterns. In this case, query, stream_query, bidi_stream_query.\n",
        "    *  query(...): The classic, synchronous request-response pattern. It takes one input and returns one complete output.\n",
        "    *  stream_query(...): The server-side streaming pattern. It takes one input but returns a Python generator (yield), allowing it to stream multiple chunks back to the client. This is perfect for streaming LLM text responses.\n",
        "    *  bidi_stream_query(...): The new full bidirectional streaming pattern. This is an async method that receives an asyncio.Queue. It runs in a while True loop, waiting for messages to arrive on the queue (await queue.get()) and using yield to send responses back at any time. This enables true two-way conversation and is the pattern the Live API is built on.\n",
        "\n",
        "Finally the register_operations(...) is the agent's \"routing table.\" It tells Agent Engine which class method to call for which endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dVOvLbuJq3a4"
      },
      "source": [
        "### Create the Echo Agent\n",
        "\n",
        "Let's implement a simple echo agent that demonstrates these interaction patterns supported by Agent Engine.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oZGXka3vq5rn"
      },
      "outputs": [],
      "source": [
        "class EchoAgent:\n",
        "    \"\"\"\n",
        "    A simple echo agent demonstrating Agent Engine patterns.\n",
        "\n",
        "    This agent showcases three interaction patterns:\n",
        "    1. Standard query (request-response)\n",
        "    2. Streaming query (server-side streaming)\n",
        "    3. Bidirectional streaming (two-way communication)\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, project: str, location: str):\n",
        "        \"\"\"\n",
        "        Initialize configuration (keep lightweight and pickle-able).\n",
        "\n",
        "        Important: Only store simple data types here.\n",
        "        Heavy initialization should go in set_up().\n",
        "        \"\"\"\n",
        "        self.project = project\n",
        "        self.location = location\n",
        "\n",
        "    def set_up(self):\n",
        "        \"\"\"\n",
        "        Perform heavy initialization (called once when deployed).\n",
        "\n",
        "        This is where you'd initialize:\n",
        "        - Database connections\n",
        "        - Model loading\n",
        "        - External service clients\n",
        "        \"\"\"\n",
        "        import logging\n",
        "        self.logger = logging.getLogger(__name__)\n",
        "        self.logger.info(\"Echo Agent ready!\")\n",
        "\n",
        "    def query(self, input: str) -> Dict[str, Any]:\n",
        "        \"\"\"\n",
        "        Standard query - returns complete response.\n",
        "\n",
        "        This is the simplest pattern: receive input, return output.\n",
        "        Use this for quick, synchronous operations.\n",
        "        \"\"\"\n",
        "        return {\n",
        "            \"input\": input,\n",
        "            \"output\": f\"Echo: {input}\",\n",
        "            \"method\": \"query\"\n",
        "        }\n",
        "\n",
        "    def stream_query(self, input: str) -> Iterator[Dict[str, Any]]:\n",
        "        \"\"\"\n",
        "        Streaming query - yields response chunks.\n",
        "\n",
        "        Use this pattern when you want to:\n",
        "        - Stream partial results as they become available\n",
        "        - Handle large responses that shouldn't be buffered\n",
        "        - Provide progress updates during long operations\n",
        "        \"\"\"\n",
        "        import time\n",
        "\n",
        "        words = input.split()\n",
        "\n",
        "        # Agent Engine expects each yield to be a complete, serializable response\n",
        "        for i, word in enumerate(words):\n",
        "            # Each chunk should be a complete response object\n",
        "            chunk_response = {\n",
        "                \"chunk\": f\"{word} \",\n",
        "                \"progress\": f\"{i+1}/{len(words)}\"\n",
        "            }\n",
        "\n",
        "            yield chunk_response\n",
        "\n",
        "            # Simulate processing time\n",
        "            await asyncio.sleep(0.1)\n",
        "\n",
        "    async def bidi_stream_query(\n",
        "        self,\n",
        "        queue: asyncio.Queue\n",
        "    ) -> AsyncIterator[Dict[str, Any]]:\n",
        "        \"\"\"\n",
        "        Bidirectional streaming for continuous conversation.\n",
        "\n",
        "        This enables:\n",
        "        - Real-time, two-way communication\n",
        "        - User interruptions mid-response\n",
        "        - Long-running interactive sessions\n",
        "        \"\"\"\n",
        "        self.logger.info(\"Bidi session started\")\n",
        "\n",
        "        while True:\n",
        "            # Wait for message from the queue\n",
        "            message = await queue.get()\n",
        "            user_input = message.get(\"input\", \"\")\n",
        "\n",
        "            # Check for exit command\n",
        "            if user_input.lower() in (\"exit\", \"quit\"):\n",
        "                yield {\"output\": \"Goodbye!\"}\n",
        "                break\n",
        "\n",
        "            # Echo back the input\n",
        "            # In a real agent, this is where you'd process the input\n",
        "            yield {\"output\": f\"Echo: {user_input}\"}\n",
        "\n",
        "    def register_operations(self):\n",
        "        \"\"\"\n",
        "        Register available operations with Agent Engine.\n",
        "\n",
        "        This tells Agent Engine which methods are available\n",
        "        and how to route requests to them.\n",
        "        \"\"\"\n",
        "        return {\n",
        "            \"\": [\"query\"],                        # Default endpoint\n",
        "            \"stream\": [\"stream_query\"],           # Streaming endpoint\n",
        "            \"bidi_stream\": [\"bidi_stream_query\"]  # Bidirectional endpoint\n",
        "        }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OMR_RN0-zYFi"
      },
      "source": [
        "### Deploy the Echo Agent\n",
        "\n",
        "Now let's deploy our echo agent to Agent Engine.\n",
        "\n",
        "First, we create a local instance of our EchoAgent class.\n",
        "\n",
        "Then, we call client.agent_engines.create(). This single command is the core of Agent Engine. It takes our local Python object (agent=echo_agent), serializes it, determines its dependencies (which we specify in requirements), stages all these files in our Cloud Storage, provisions the necessary serverless infrastructure, and deploys our agent as a scalable, secure endpoint.\n",
        "\n",
        "Once deployed, you have the remote_echo_agent which is a client object that we can use to interact with our newly deployed agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "knsbVEy4zbQy"
      },
      "outputs": [],
      "source": [
        "# Create an instance of our echo agent\n",
        "echo_agent = EchoAgent(\n",
        "    project=PROJECT_ID,\n",
        "    location=LOCATION\n",
        ")\n",
        "\n",
        "# Deploy to Agent Engine\n",
        "remote_echo_agent = client.agent_engines.create(\n",
        "    agent=echo_agent,\n",
        "    config={\n",
        "        # Display name in the console\n",
        "        \"display_name\": \"Echo Agent Tutorial\",\n",
        "        # Description for documentation\n",
        "        \"description\": \"Simple echo agent demonstrating streaming patterns\",\n",
        "        # Python dependencies needed by the agent\n",
        "        \"requirements\": [\n",
        "            \"google-cloud-aiplatform[agent_engines] @ git+https://github.com/googleapis/python-aiplatform.git\"\n",
        "        ],\n",
        "        # Cloud Storage bucket for staging\n",
        "        \"staging_bucket\": BUCKET_URI,\n",
        "        # Set the experimental mode\n",
        "        \"agent_server_mode\": vertexai_types.AgentServerMode.EXPERIMENTAL\n",
        "    }\n",
        ")\n",
        "\n",
        "print(\"✅ Agent deployed successfully!\")\n",
        "print(f\"Resource name: {remote_echo_agent.api_resource.name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sdk6QLPDvm7X"
      },
      "source": [
        "### Test the agent\n",
        "\n",
        "Let's test all three interaction patterns to verify our agent is working correctly.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WR2peiCOwa3r"
      },
      "source": [
        "#### Test Standard Query\n",
        "\n",
        "Here, we test the simplest pattern. We call the .query() method on our remote agent handle. The SDK sends the request, waits for the complete JSON response, and then prints it.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PgnM2iFrweBi"
      },
      "outputs": [],
      "source": [
        "# Test 1: Standard query (request-response)\n",
        "print(\"=\" * 50)\n",
        "print(\"Testing standard query()...\")\n",
        "print(\"=\" * 50)\n",
        "\n",
        "try:\n",
        "    response = remote_echo_agent.query(input=\"Hello from standard query\")\n",
        "    print(f\"✅ Query response: {response}\")\n",
        "except Exception as e:\n",
        "    print(f\"❌ Query failed: {e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8PamPWGPwiJq"
      },
      "source": [
        "#### Test Streaming Query\n",
        "\n",
        "This tests the server-side streaming pattern. We call .stream_query() on the remote agent. Because the underlying method in our class is a generator, the client object also behaves like a generator. This lets us use a simple for loop to iterate over the response chunks as they're streamed back from the server, printing each one as it arrives.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bn4qEuE5wmdP"
      },
      "outputs": [],
      "source": [
        "# Test 2: Streaming query (server-side streaming)\n",
        "print(\"\\n\" + \"=\" * 50)\n",
        "print(\"Testing stream_query()...\")\n",
        "print(\"=\" * 50)\n",
        "\n",
        "try:\n",
        "    print(\"Streaming response:\")\n",
        "    for chunk in remote_echo_agent.stream_query(input=\"Hello from streaming query\"):\n",
        "        print(f\"  Received chunk: {chunk}\")\n",
        "    print(\"✅ Streaming completed\")\n",
        "except Exception as e:\n",
        "    print(f\"❌ Stream query failed: {e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Up7lyx48wr6G"
      },
      "source": [
        "#### Test Bidirectional Streaming\n",
        "\n",
        "This tests our most advanced pattern. Because it's asynchronous, we must define an async test function.\n",
        "\n",
        "*   We use async with client.aio.live.agent_engines.connect(...) to open a persistent, two-way connection to the agent endpoint we routed to bidi_stream_query.\n",
        "*   This gives us a session object with two primary methods: session.send() (to send data to the agent) and session.receive() (to get data from the agent).\n",
        "*   We simulate a conversation: we .send() a message, await its echoed .receive() response, and repeat.\n",
        "*   Finally, we send \"exit,\" which our agent's logic is designed to catch, causing it to yield a final \"Goodbye!\" and break its loop.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L-C0BQUhxaJG"
      },
      "outputs": [],
      "source": [
        "# Test 3: Bidirectional streaming (two-way communication)\n",
        "print(\"\\n\" + \"=\" * 50)\n",
        "print(\"Testing bidi_stream_query()...\")\n",
        "print(\"=\" * 50)\n",
        "\n",
        "async def test_bidi_streaming():\n",
        "    \"\"\"Test bidirectional streaming with multiple exchanges.\"\"\"\n",
        "\n",
        "    try:\n",
        "        # Connect to the agent's bidirectional streaming endpoint\n",
        "        async with client.aio.live.agent_engines.connect(\n",
        "            agent_engine=remote_echo_agent.api_resource.name,\n",
        "            config={\"class_method\": \"bidi_stream_query\"}\n",
        "        ) as session:\n",
        "\n",
        "            # Exchange 1: Send greeting\n",
        "            await session.send({\"input\": \"Hello from bidi\"})\n",
        "            response = await session.receive()\n",
        "            print(\"Exchange 1 - Sent: 'Hello from bidi'\")\n",
        "            print(f\"Exchange 1 - Received: {response['bidiStreamOutput']['output']}\")\n",
        "\n",
        "            # Exchange 2: Send another message\n",
        "            await session.send({\"input\": \"How are you?\"})\n",
        "            response = await session.receive()\n",
        "            print(\"Exchange 2 - Sent: 'How are you?'\")\n",
        "            print(f\"Exchange 2 - Received: {response['bidiStreamOutput']['output']}\")\n",
        "\n",
        "            # Exchange 3: Exit\n",
        "            await session.send({\"input\": \"exit\"})\n",
        "            response = await session.receive()\n",
        "            print(\"Exchange 3 - Sent: 'exit'\")\n",
        "            print(f\"Exchange 3 - Received: {response['bidiStreamOutput']['output']}\")\n",
        "\n",
        "            await session.close()\n",
        "            print(\"✅ Bidirectional streaming completed\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"❌ Bidi stream failed: {e}\")\n",
        "\n",
        "# Run the async test\n",
        "await test_bidi_streaming()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R1uWu8lMLAV_"
      },
      "source": [
        "## Live API Audio Conversations\n",
        "\n",
        "Now that we understand the basic patterns, let's explore the Live API's real-time audio capabilities. This enables natural voice conversations with your agents.\n",
        "\n",
        "### Understanding the Live API\n",
        "\n",
        "The Live API is Google's real-time interface to Gemini models, enabling:\n",
        "\n",
        "- **WebSocket connection**: Direct, persistent connection to Gemini models\n",
        "- **Multimodal streaming**: Support for text, audio, and video in real-time\n",
        "- **Natural interruptions**: Users can interrupt the agent mid-response\n",
        "- **Low latency**: Optimized for conversational experiences (< 1 second response time)\n",
        "- **Session continuity**: Maintains context throughout the conversation\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HywFxuqUN8FI"
      },
      "source": [
        "### Create an agent using Live API\n",
        "\n",
        "This agent connects directly to the Gemini Live API via WebSockets to enable a real-time audio conversation.\n",
        "\n",
        "In this case, we have the set_up(...) to build the specific WebSocket URL (wss://...) for the Live API and set the generation config to explicitly request AUDIO responses.\n",
        "\n",
        "Then we have the bidi_stream_query(...) which is the main entry point that Agent Engine calls:\n",
        "\n",
        "*   It gets an auth token using _authenticate().\n",
        "*   It opens a WebSocket connection to the Live API using the websockets library.\n",
        "*   It initializes the connection using _setup_session().\n",
        "*   It enters its main loop: it waits for text from our notebook client (await input_queue.get()), sends that text to Gemini with _send_text(), and then streams the audio response back by iterating over _receive_audio(), yielding each audio chunk back to the notebook client.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kPDlkfEsLEOi"
      },
      "outputs": [],
      "source": [
        "class LiveAudioAgent:\n",
        "    \"\"\"\n",
        "    Agent with real-time audio capabilities using Live API.\n",
        "\n",
        "    This implementation demonstrates:\n",
        "    - WebSocket connection to Gemini Live API\n",
        "    - Audio streaming configuration\n",
        "    - Session management\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, project: str, location: str, model_id: str):\n",
        "        \"\"\"Initialize with minimal configuration.\"\"\"\n",
        "        self.project = project\n",
        "        self.location = location\n",
        "        self.model_id = model_id\n",
        "\n",
        "    def set_up(self):\n",
        "        \"\"\"Configure Live API connection.\"\"\"\n",
        "        # Build WebSocket URL for Live API\n",
        "        host = f\"{self.location}-aiplatform.googleapis.com\"\n",
        "        self.service_url = (\n",
        "            f\"wss://{host}/ws/google.cloud.aiplatform.v1.\"\n",
        "            \"LlmBidiService/BidiGenerateContent\"\n",
        "        )\n",
        "\n",
        "        # Model path for Vertex AI\n",
        "        self.model = (\n",
        "            f\"projects/{self.project}/locations/{self.location}/\"\n",
        "            f\"publishers/google/models/{self.model_id}\"\n",
        "        )\n",
        "\n",
        "        # Configure for audio responses\n",
        "        self.config = {\n",
        "            \"response_modalities\": [\"AUDIO\"]\n",
        "        }\n",
        "\n",
        "    async def _authenticate(self):\n",
        "        \"\"\"Get authentication token for WebSocket connection.\"\"\"\n",
        "        credentials, _ = google.auth.default()\n",
        "        auth_req = google.auth.transport.requests.Request()\n",
        "        credentials.refresh(auth_req)\n",
        "        return credentials.token\n",
        "\n",
        "    async def _setup_session(self, websocket):\n",
        "        \"\"\"Initialize Live API session.\"\"\"\n",
        "        setup_message = {\n",
        "            \"setup\": {\n",
        "                \"model\": self.model,\n",
        "                \"generation_config\": self.config,\n",
        "            }\n",
        "        }\n",
        "        await websocket.send(json.dumps(setup_message))\n",
        "\n",
        "        # Receive setup confirmation\n",
        "        response = await websocket.recv(decode=False)\n",
        "        setup_response = json.loads(response.decode(\"ascii\"))\n",
        "        logging.info(f\"Session established: {setup_response}\")\n",
        "\n",
        "    async def _send_text(self, websocket, text: str) -> bool:\n",
        "        \"\"\"\n",
        "        Send text message to Live API.\n",
        "\n",
        "        Returns:\n",
        "            False if exit requested, True otherwise\n",
        "        \"\"\"\n",
        "        if text.lower() in (\"exit\", \"quit\"):\n",
        "            return False\n",
        "\n",
        "        message = {\n",
        "            \"client_content\": {\n",
        "                \"turns\": [{\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [{\"text\": text}]\n",
        "                }],\n",
        "                \"turn_complete\": True,\n",
        "            }\n",
        "        }\n",
        "        await websocket.send(json.dumps(message))\n",
        "        return True\n",
        "\n",
        "    async def _receive_audio(self, websocket) -> AsyncIterator[Dict]:\n",
        "        \"\"\"\n",
        "        Receive and process audio response from Live API.\n",
        "\n",
        "        Yields:\n",
        "            Audio data dictionaries\n",
        "        \"\"\"\n",
        "        async for raw_response in websocket:\n",
        "            response = json.loads(raw_response.decode())\n",
        "            server_content = response.get(\"serverContent\")\n",
        "\n",
        "            if not server_content:\n",
        "                break\n",
        "\n",
        "            # Extract audio from model turn\n",
        "            model_turn = server_content.get(\"modelTurn\")\n",
        "            if model_turn:\n",
        "                parts = model_turn.get(\"parts\", [])\n",
        "                for part in parts:\n",
        "                    if \"inlineData\" in part:\n",
        "                        # Decode PCM audio data\n",
        "                        pcm_data = base64.b64decode(\n",
        "                            part[\"inlineData\"][\"data\"]\n",
        "                        )\n",
        "                        audio_array = np.frombuffer(\n",
        "                            pcm_data,\n",
        "                            dtype=np.int16\n",
        "                        )\n",
        "                        yield {\"output\": audio_array.tolist()}\n",
        "\n",
        "            # Check for turn completion\n",
        "            if server_content.get(\"turnComplete\"):\n",
        "                break\n",
        "\n",
        "        yield {\"output\": \"end of turn\"}\n",
        "\n",
        "    async def bidi_stream_query(\n",
        "        self,\n",
        "        input_queue: asyncio.Queue\n",
        "    ) -> AsyncIterator[Dict[str, Any]]:\n",
        "        \"\"\"\n",
        "        Handle bidirectional audio streaming.\n",
        "\n",
        "        Args:\n",
        "            input_queue: Queue for receiving user input\n",
        "\n",
        "        Yields:\n",
        "            Audio response data\n",
        "        \"\"\"\n",
        "        # Get authentication token\n",
        "        access_token = await self._authenticate()\n",
        "\n",
        "        headers = {\n",
        "            \"Content-Type\": \"application/json\",\n",
        "            \"Authorization\": f\"Bearer {access_token}\",\n",
        "        }\n",
        "\n",
        "        # Establish WebSocket connection\n",
        "        async with websockets.asyncio.client.connect(\n",
        "            self.service_url,\n",
        "            additional_headers=headers\n",
        "        ) as websocket:\n",
        "\n",
        "            # Initialize session\n",
        "            await self._setup_session(websocket)\n",
        "\n",
        "            # Main conversation loop\n",
        "            while True:\n",
        "                # Get user input\n",
        "                request = await input_queue.get()\n",
        "                text_input = request[\"input\"]\n",
        "\n",
        "                # Send to Live API\n",
        "                if not await self._send_text(websocket, text_input):\n",
        "                    break\n",
        "\n",
        "                # Receive and yield audio response\n",
        "                async for audio_data in self._receive_audio(websocket):\n",
        "                    yield audio_data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ncE5aDRxN_xF"
      },
      "source": [
        "### Deploy the Live Audio Agent\n",
        "\n",
        "Let's deploy the agent. The process is identical to the first agent, but notice the requirements section in the config. We must list numpy, google-auth, and websockets because our LiveAudioAgent class imports and uses them. Agent Engine will see this list and pip install these packages into the agent's serverless container."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1fTy-gp4OCRH"
      },
      "outputs": [],
      "source": [
        "# Create Live Audio Agent\n",
        "live_agent = LiveAudioAgent(\n",
        "    project=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    model_id=MODEL_ID  # Use a model that supports live API\n",
        ")\n",
        "\n",
        "# Deploy agent\n",
        "remote_live_agent = client.agent_engines.create(\n",
        "    agent=live_agent,\n",
        "    config={\n",
        "        \"display_name\": \"Live Audio Agent Tutorial\",\n",
        "        \"description\": \"Simple agent with real-time audio capabilities\",\n",
        "        \"requirements\": [\n",
        "            \"numpy\",\n",
        "            \"google-auth\",\n",
        "            \"websockets\"\n",
        "        ],\n",
        "        \"staging_bucket\": BUCKET_URI,\n",
        "        # Set the experimental mode\n",
        "        \"agent_server_mode\": vertexai_types.AgentServerMode.EXPERIMENTAL\n",
        "    },\n",
        ")\n",
        "\n",
        "print(\"✅ Agent deployed successfully!\")\n",
        "print(f\"Resource name: {remote_live_agent.api_resource.name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ixtxx32JRfoh"
      },
      "source": [
        "### Test the live audio agent\n",
        "\n",
        "To interact with the audio agent, you define a simple test client.\n",
        "\n",
        "The client connects to our deployed agent's bidi_stream_query endpoint using the familiar client.aio.live.agent_engines.connect pattern.\n",
        "\n",
        "It enters a while True loop that takes text from the console using input(\"You: \"). It sends this text to our deployed agent using session.send().\n",
        "\n",
        "It then enters an inner loop to collect the response. It calls await session.receive() repeatedly, appending each received audio array (which our agent is yielding) to the audio_chunks list. This inner loop breaks when it receives the special \"end of turn\" string that our agent yields when Gemini is done speaking.\n",
        "\n",
        "Finally, it uses np.concatenate to stitch all the small audio chunks into one complete audio stream and uses display(Audio(...)) to play it directly in the notebook, creating a seamless text-in, audio-out chat experience.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PkZ7Dg14RZnq"
      },
      "outputs": [],
      "source": [
        "async def interactive_audio_chat(client, agent_name: str):\n",
        "    \"\"\"Simple audio chat interface.\"\"\"\n",
        "\n",
        "    async with client.aio.live.agent_engines.connect(\n",
        "        agent_engine=agent_name,\n",
        "        config={\"class_method\": \"bidi_stream_query\"}\n",
        "    ) as session:\n",
        "\n",
        "        print(\"🎤 Audio Chat Ready! (type 'exit' to quit)\\n\")\n",
        "\n",
        "        while True:\n",
        "            await asyncio.sleep(0.1)\n",
        "            user_input = input(\"You: \")\n",
        "\n",
        "            if user_input.lower() == \"exit\":\n",
        "                await session.send({\"input\": user_input})\n",
        "                break\n",
        "\n",
        "            # Send text\n",
        "            await session.send({\"input\": user_input})\n",
        "\n",
        "            # Collect audio response\n",
        "            audio_chunks = []\n",
        "\n",
        "            while True:\n",
        "                response = await session.receive()\n",
        "                output = response[\"bidiStreamOutput\"]\n",
        "                if output.get(\"output\") == \"end of turn\":\n",
        "                  break\n",
        "                audio = output[\"output\"]\n",
        "                audio_chunks.append(np.array(audio))\n",
        "\n",
        "            # Play audio\n",
        "            print(\"Agent:\")\n",
        "            if audio_chunks:\n",
        "                full_audio = np.concatenate(audio_chunks)\n",
        "                display(Audio(full_audio, rate=24000, autoplay=True))\n",
        "                print()\n",
        "\n",
        "        await session.close()\n",
        "\n",
        "# Run audio chat\n",
        "await interactive_audio_chat(client, remote_live_agent.api_resource.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0u5wOvlXiUkW"
      },
      "source": [
        "## Creating a Weather Assistant with ADK\n",
        "\n",
        "In this final part, we'll build a weather assistant using the Agent Development Kit (ADK). This demonstrates how to combine Live API capabilities with tools, session management, and memory services.\n",
        "\n",
        "### Understanding ADK Integration\n",
        "\n",
        "The Agent Development Kit (ADK) provides high-level abstractions for building sophisticated agents:\n",
        "\n",
        "- **Tool integration**: Easy function calling with automatic orchestration\n",
        "- **Session management**: Persistent user context across conversations\n",
        "- **Memory services**: Automatic conversation history management\n",
        "- **Built-in patterns**: Production best practices by default\n",
        "- **Error handling**: Robust error recovery and retry logic\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W-jxqWcy3JXC"
      },
      "source": [
        "### Define simple weather tools\n",
        "\n",
        "This is the first tool for our ADK agent. It takes location (which the LLM will provide) and a special tool_context argument and returns a dictionary of (mock) weather data. The ADK will automatically serialize this and send it back to the LLM.\n",
        "\n",
        "Notice the tool uses the ToolContext, which is an ADK object that gives the tool access to a persistent state. Here, we use tool_context.state (a dictionary) to store a history of all weather queries, demonstrating how tools can manage state."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Og5rTztk2tut"
      },
      "outputs": [],
      "source": [
        "async def get_weather(\n",
        "    location: str,\n",
        "    tool_context: ToolContext\n",
        ") -> Dict[str, str]:\n",
        "    \"\"\"\n",
        "    Get current weather for a location.\n",
        "\n",
        "    This is a simple mock implementation.\n",
        "    In production, integrate with a real weather API.\n",
        "\n",
        "    Args:\n",
        "        location: City or location name\n",
        "        tool_context: ADK tool context for state management\n",
        "\n",
        "    Returns:\n",
        "        Weather information dictionary\n",
        "    \"\"\"\n",
        "\n",
        "    import random\n",
        "    from datetime import datetime\n",
        "\n",
        "    # Mock weather data (replace with real API)\n",
        "    temperature_c = random.randint(10, 35)\n",
        "    temperature_f = int(temperature_c * 9/5 + 32)\n",
        "    conditions = random.choice([\n",
        "        \"Sunny\", \"Cloudy\", \"Rainy\", \"Partly Cloudy\"\n",
        "    ])\n",
        "\n",
        "    # Store in tool context for history\n",
        "    if \"weather_queries\" not in tool_context.state:\n",
        "        tool_context.state[\"weather_queries\"] = []\n",
        "\n",
        "    tool_context.state[\"weather_queries\"].append({\n",
        "        \"location\": location,\n",
        "        \"timestamp\": datetime.now().isoformat(),\n",
        "        \"temperature\": f\"{temperature_c}°C\"\n",
        "    })\n",
        "\n",
        "    return {\n",
        "        \"location\": location,\n",
        "        \"temperature\": f\"{temperature_c}°C / {temperature_f}°F\",\n",
        "        \"conditions\": conditions,\n",
        "        \"humidity\": f\"{random.randint(30, 80)}%\",\n",
        "        \"wind\": f\"{random.randint(5, 25)} km/h\"\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TjhGWeJF9hiN"
      },
      "source": [
        "This is our second tool. It's a simpler function that just takes a location and number of days and returns a JSON string of a (mock) forecast. It doesn't need the tool_context since it's not saving any state."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "idoNume03SMe"
      },
      "outputs": [],
      "source": [
        "async def get_forecast(\n",
        "    location: str,\n",
        "    days: int = 3\n",
        ") -> str:\n",
        "    \"\"\"\n",
        "    Get weather forecast for multiple days.\n",
        "\n",
        "    Args:\n",
        "        location: City or location name\n",
        "        days: Number of days to forecast (max 5)\n",
        "\n",
        "    Returns:\n",
        "        JSON string with forecast data\n",
        "    \"\"\"\n",
        "    import random\n",
        "    from datetime import datetime, timedelta\n",
        "    import json\n",
        "\n",
        "    days = min(days, 5)  # Limit to 5 days\n",
        "    forecast = []\n",
        "\n",
        "    for i in range(days):\n",
        "        date = datetime.now() + timedelta(days=i)\n",
        "        forecast.append({\n",
        "            \"date\": date.strftime(\"%A, %B %d\"),\n",
        "            \"high\": f\"{random.randint(15, 30)}°C\",\n",
        "            \"low\": f\"{random.randint(5, 20)}°C\",\n",
        "            \"conditions\": random.choice([\n",
        "                \"Sunny\", \"Cloudy\", \"Rainy\", \"Partly Cloudy\"\n",
        "            ])\n",
        "        })\n",
        "\n",
        "    return json.dumps(forecast, indent=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W2YlorTS4eWO"
      },
      "source": [
        "### Define the ADK agent\n",
        "\n",
        "Instead of manually writing a complex LiveAudioAgent class, we use the high-level LlmAgent class from the ADK.\n",
        "\n",
        "We specify the model (which must support the Live API) and a detailed instruction (system prompt). And, we pass our two Python functions directly into the tools=[get_weather, get_forecast] list.\n",
        "\n",
        "The ADK takes care of everything else: using this class together with other primitives, it will inspect the functions to create the tool definition (function calling schema), orchestrate the LLM calls, parse the LLM's request to use a tool, execute our Python function, and send the result back to the LLM to generate a final answer.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l418opOT4gjo"
      },
      "outputs": [],
      "source": [
        "# Create the ADK Weather Agent\n",
        "weather_agent = LlmAgent(\n",
        "        model='gemini-2.0-flash-live-preview-04-09',\n",
        "        name='weather_assistant',\n",
        "        description='A helpful weather assistant with current conditions and forecasts',\n",
        "        instruction=\"\"\"\n",
        "        You are a friendly and knowledgeable weather assistant.\n",
        "\n",
        "        Your responsibilities:\n",
        "        1. Provide current weather when asked\n",
        "        2. Give forecasts when requested\n",
        "        3. Be conversational and helpful\n",
        "        4. Always mention the location in responses\n",
        "        5. Provide temperature in both Celsius and Fahrenheit\n",
        "        6. If asked for a forecast, provide all requested days\n",
        "\n",
        "        Guidelines:\n",
        "        - Use the get_weather tool for current conditions\n",
        "        - Use the get_forecast tool for multi-day forecasts\n",
        "        - Be proactive in offering additional information\n",
        "        - Keep responses concise but informative\n",
        "        \"\"\",\n",
        "        tools=[\n",
        "            get_weather,\n",
        "            get_forecast,\n",
        "        ],\n",
        "        generate_content_config=types.GenerateContentConfig(\n",
        "            temperature=0.7,\n",
        "            max_output_tokens=512,\n",
        "            safety_settings=[\n",
        "                types.SafetySetting(\n",
        "                    category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,\n",
        "                    threshold=types.HarmBlockThreshold.OFF,\n",
        "                )\n",
        "            ]\n",
        "        )\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lXRui40u4zBD"
      },
      "source": [
        "### Deploy the agent on Agent Engine\n",
        "\n",
        "To be deployed on Agent Engine, an ADK Agent must be wrapped in an AdkApp. This wrapper provides all the necessary methods (like bidi_stream_query) that Agent Engine expects.\n",
        "\n",
        "In this scenario, we provide builders for InMemorySessionService and InMemoryMemoryService. This tells the ADK to automatically manage conversation history and session context for us (in this case, just in memory). For production, you could swap these with persistent services like Firestore.\n",
        "\n",
        "Finally, we call client.agent_engines.create() as before, but this time passing agent=app. We also add google-adk and google-genai to the requirements list."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9jGJdVu341WX"
      },
      "outputs": [],
      "source": [
        "# Set up session and memory services\n",
        "def session_service_builder():\n",
        "    from google.adk.sessions.in_memory_session_service import InMemorySessionService\n",
        "    return InMemorySessionService()\n",
        "\n",
        "def memory_service_builder():\n",
        "    from google.adk.memory.in_memory_memory_service import InMemoryMemoryService\n",
        "    return InMemoryMemoryService()\n",
        "    \n",
        "# Wrap agent with AdkApp - this is all you need!\n",
        "app = AdkApp(agent=weather_agent, session_service_builder=session_service_builder, memory_service_builder=memory_service_builder)\n",
        "\n",
        "remote_live_adk_agent = client.agent_engines.create(\n",
        "    agent=app,\n",
        "    config={\n",
        "        \"display_name\": \"Weather Assistant Tutorial\",\n",
        "        \"description\": \"Simple agent with real-time audio capabilities\",\n",
        "        \"requirements\": [\n",
        "            \"google-adk\",\n",
        "            \"google-genai\",\n",
        "            \"google-cloud-aiplatform[agent_engines] @ git+https://github.com/googleapis/python-aiplatform.git\",\n",
        "        ],\n",
        "        \"staging_bucket\": BUCKET_URI,\n",
        "        # Set the experimental mode\n",
        "        \"agent_server_mode\": vertexai_types.AgentServerMode.EXPERIMENTAL\n",
        "    },\n",
        ")\n",
        "\n",
        "print(\"✅ Agent deployed successfully!\")\n",
        "print(f\"Resource name: {remote_live_adk_agent.api_resource.name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aqSAPLrH5kOz"
      },
      "source": [
        "### Test the ADK agent\n",
        "\n",
        "This is our final test client, built to interact with the deployed ADK agent.\n",
        "\n",
        "The client defines a helper, prepare_live_request, to wrap our text input into the specific LiveRequest object that the deployed AdkApp expects. We connect using the same async with client.aio.live.agent_engines.connect pattern.\n",
        "\n",
        "The main difference is in the response-handling loop. The ADK streams back a structured Event object. This event can contain text, tool-use information, or audio. Our inner loop parses this Event. If it contains inline_data, it's an audio chunk, which we collect just like in our previous agent. If it contains part.text, it's a text update (which we print).This enables the agent to provide both audio and text responses, managed automatically by the ADK.\n",
        "\n",
        "When all events for the turn are received, we concatenate and play the full audio."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kd5m0UOoyXxU"
      },
      "outputs": [],
      "source": [
        "async def weather_chat_with_audio(client, agent_resource_name: str):\n",
        "    \"\"\"Interactive chat with the deployed weather agent.\"\"\"\n",
        "\n",
        "    def prepare_live_request(input_text: str):\n",
        "        \"\"\"Prepare request for ADK agent.\"\"\"\n",
        "        part = types.Part.from_text(text=input_text)\n",
        "        content = types.Content(parts=[part])\n",
        "        return LiveRequest(content=content)\n",
        "\n",
        "    async with client.aio.live.agent_engines.connect(\n",
        "        agent_engine=agent_resource_name,\n",
        "        config={\"class_method\": \"bidi_stream_query\"}\n",
        "    ) as connection:\n",
        "\n",
        "        print(\"🌤️ Weather Assistant Ready!\")\n",
        "        print(\"Ask about weather in any city, or type 'exit' to quit.\\n\")\n",
        "        print(\"Example questions:\")\n",
        "        print(\"- What's the weather in San Francisco?\")\n",
        "        print(\"- Give me a 5-day forecast for Tokyo\")\n",
        "        print(\"- Is it raining in London?\\n\")\n",
        "\n",
        "        user_id = \"weather_user\"\n",
        "        first_req = True\n",
        "\n",
        "        while True:\n",
        "            # Get user input\n",
        "            time.sleep(0.1)  # Small delay for display\n",
        "            input_text = input(\"You: \")\n",
        "\n",
        "            if input_text.lower() == \"exit\":\n",
        "                break\n",
        "\n",
        "            # Prepare for audio response\n",
        "            audio_data = []\n",
        "            seen_audio = False\n",
        "\n",
        "            # Send request\n",
        "            if first_req:\n",
        "                await connection.send({\n",
        "                    \"user_id\": user_id,\n",
        "                    \"live_request\": prepare_live_request(input_text).dict()\n",
        "                })\n",
        "                first_req = False\n",
        "            else:\n",
        "                await connection.send(prepare_live_request(input_text).dict())\n",
        "\n",
        "            # Receive response\n",
        "            print(\"Weather Assistant: \", end=\"\")\n",
        "\n",
        "            while True:\n",
        "                received = await connection.receive()\n",
        "                event = received[\"bidiStreamOutput\"]\n",
        "\n",
        "                # Handle the response format\n",
        "                if \"actions\" in event and \"requested_tool_confirmations\" in event[\"actions\"]:\n",
        "                    del event[\"actions\"][\"requested_tool_confirmations\"]\n",
        "\n",
        "                event = Event.model_validate(event)\n",
        "\n",
        "                # Extract content\n",
        "                part = (\n",
        "                    event.content and event.content.parts and event.content.parts[0]\n",
        "                )\n",
        "\n",
        "                if not part and seen_audio:\n",
        "                    break\n",
        "\n",
        "                if part and part.inline_data and part.inline_data.data:\n",
        "                    # Audio response\n",
        "                    seen_audio = True\n",
        "                    chunk_data = part.inline_data.data\n",
        "                    data = np.frombuffer(chunk_data, dtype=np.int16)\n",
        "                    audio_data.append(data)\n",
        "                elif part and part.text:\n",
        "                    # Text response\n",
        "                    print(part.text)\n",
        "                    break\n",
        "                elif part:\n",
        "                    # Other response types\n",
        "                    print(part)\n",
        "                else:\n",
        "                    # No more content\n",
        "                    break\n",
        "\n",
        "            # Play audio if available\n",
        "            if audio_data:\n",
        "                concatenated_audio = np.concatenate(audio_data)\n",
        "                display(Audio(concatenated_audio, rate=24000, autoplay=True))\n",
        "                print(\"[Audio response played]\")\n",
        "\n",
        "            print()  # New line for next interaction\n",
        "\n",
        "        print(\"\\n👋 Thanks for using Weather Assistant!\")\n",
        "\n",
        "# Run audio chat\n",
        "await weather_chat_with_audio(client, remote_live_adk_agent.api_resource.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To avoid incurring charges for the resources created in this tutorial, clean up the deployed agents."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0SD9rlW_u59x"
      },
      "outputs": [],
      "source": [
        "# Set to True to delete the agents\n",
        "delete_agents = True  # @param {type: \"boolean\"}\n",
        "\n",
        "if delete_agents:\n",
        "    try:\n",
        "        # Delete the echo agent\n",
        "        print(\"Deleting Echo Agent...\")\n",
        "        remote_echo_agent.delete(force=True)\n",
        "        print(\"✅ Echo Agent deleted\")\n",
        "    except Exception as e:\n",
        "        print(f\"❌ Failed to delete Echo Agent: {e}\")\n",
        "\n",
        "    try:\n",
        "        # Delete the live audio agent\n",
        "        print(\"Deleting Live Audio Agent...\")\n",
        "        remote_live_agent.delete(force=True)\n",
        "        print(\"✅ Live Audio Agent deleted\")\n",
        "    except Exception as e:\n",
        "        print(f\"❌ Failed to delete Live Audio Agent: {e}\")\n",
        "\n",
        "    try:\n",
        "        # Delete the weather assistant\n",
        "        print(\"Deleting Weather Assistant...\")\n",
        "        remote_live_adk_agent.delete(force=True)\n",
        "        print(\"✅ Weather Assistant deleted\")\n",
        "    except Exception as e:\n",
        "        print(f\"❌ Failed to delete Weather Assistant: {e}\")\n",
        "\n",
        "    print(\"\\n✨ All resources cleaned up successfully!\")\n",
        "else:\n",
        "    print(\"⚠️ Agents not deleted. Remember to clean up resources when done.\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tutorial_get_started_with_live_api_on_agent_engine.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
