{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "id": "yWN6WGFthSVV" }, "outputs": [], "source": [ "# Copyright 2026 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": "foZ4s_XLamQc" }, "source": [ "# Getting Started with Bidirectional Streaming v2 on Agent Runtime\n", "\n", "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \"Google
Open in Colab\n", " \n", " 		\n", " \n", " \"Google
Open in Colab Enterprise\n", " \n", " 		\n", " \n", " \"Workbench
Open in Workbench\n", " \n", " 		\n", " \n", " \"GitHub
View on GitHub\n", " \n", "
\n", "\n", "
\n", "\n", "

\n", "Share to:\n", "\n", "\n", " \"LinkedIn\n", "\n", "\n", "\n", " \"Bluesky\n", "\n", "\n", "\n", " \"X\n", "\n", "\n", "\n", " \"Reddit\n", "\n", "\n", "\n", " \"Facebook\n", "\n", "

" ] }, { "cell_type": "markdown", "metadata": { "id": "qklO_oQQaeTt" }, "source": [ "| Author(s) |\n", "| --- |\n", "| Eric Gribkoff, Max Gong |" ] }, { "cell_type": "markdown", "metadata": { "id": "UmIYy1inhsYL" }, "source": [ "## Overview\n", "\n", "This tutorial demonstrates how to build, deploy, and interact with **bidirectional streaming agents** using **Agent Runtime** with Bring-Your-Own-Dockerfile Agent deployments including with the **Live API**.\n", "\n", "In this tutorial, you will:\n", "\n", "* **Build** two different types of streaming agents\n", "* **Interact** with agents\n", "* **Implement** real-time audio conversations" ] }, { "cell_type": "markdown", "metadata": { "id": "B9p8Qt9dhcl-" }, "source": [ "## Setup" ] }, { "cell_type": "markdown", "metadata": { "id": "GD2deH5OHIUL" }, "source": [ "### Define your project\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "v3NyoLyzHMMe" }, "outputs": [], "source": [ "PROJECT_ID = \"YOUR_PROJECT_ID\" # @param {type: \"string\"}" ] }, { "cell_type": "markdown", "metadata": { "id": "G0DTFsw_21K0" }, "source": [ "### Installing the Python SDK\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "hD0sLtNRq-Of" }, "outputs": [], "source": [ "%pip install --upgrade google-cloud-aiplatform[agent_engines,adk]>=1.144" ] }, { "cell_type": "markdown", "metadata": { "id": "mmR4EuDDWZEK" }, "source": [ "### Configure Authentication" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "D_d6aBmC0ZCa" }, "outputs": [], "source": [ "from google.colab import auth\n", "\n", "auth.authenticate_user()" ] }, { "cell_type": "markdown", "metadata": { "id": "FOPNlRJmWceZ" }, "source": [ "### Configure the Agent Platform client" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "yU11NBYfYgRZ" }, "outputs": [], "source": [ "import vertexai\n", "\n", "LOCATION = \"us-central1\" # @param {type: \"string\"}\n", "ENDPOINT = f\"https://{LOCATION}-aiplatform.googleapis.com\"\n", "\n", "client = vertexai.Client(\n", " project=PROJECT_ID,\n", " location=LOCATION,\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "rys_yOlaa5kE" }, "source": [ "# Echo Agent\n", "\n", "The bidirectional streaming API is compatible with arbitrary WebSocket protocols on the deployed agent. For simplicity, we start with a basic \"echo\" agent that accepts a bidirectional stream and echoes the messages back to the client. Later, we show how to use the bidirectional streaming API with a Live ADK Agent for real-time audio streaming." ] }, { "cell_type": "markdown", "metadata": { "id": "BmjeiFu3ayek" }, "source": [ "## Setting up the source files\n", "\n", "For the Echo Agent, we will deploy using Agent Runtime's bring-your-own-Docker image deployment.\n", "\n", "We define the following files and subdirectories in the current working directory:\n", "\n", "* `echo_agent/`: the AI logic for the Echo Agent. It contains the following files:\n", " * `__init__.py`\n", " * `agent.py`\n", "* `main.py`: the API server to run\n", "* `requirements.txt`: the PyPI dependencies to pick up\n", "* `Dockerfile`: the commands to assemble the image\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "y1IRh4H1Xtkd" }, "outputs": [], "source": [ "!mkdir -p echo_agent" ] }, { "cell_type": "markdown", "metadata": { "id": "qPX89vNy2-qj" }, "source": [ "### Requirements" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "0S4ALeE8XtnO" }, "outputs": [], "source": [ "%%writefile requirements.txt\n", "google-cloud-aiplatform[agent_engines,adk]\n", "fastapi\n", "uvicorn\n", "pydantic" ] }, { "cell_type": "markdown", "metadata": { "id": "sD9hT5h53BGy" }, "source": [ "### Dockerfile" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "khcDm69ds7Cp" }, "outputs": [], "source": [ "%%writefile Dockerfile\n", "FROM python:3.13-alpine\n", "WORKDIR /app\n", "\n", "RUN adduser --disabled-password --gecos \"\" myuser\n", "\n", "COPY requirements.txt .\n", "RUN pip install --no-cache-dir -r requirements.txt\n", "\n", "\n", "COPY --chown=myuser:myuser . .\n", "\n", "ENV PATH=\"/home/myuser/.local/bin:$PATH\"\n", "\n", "USER myuser\n", "\n", "CMD [\"sh\", \"-c\", \"uvicorn main:app --host 0.0.0.0 --port $PORT\"]" ] }, { "cell_type": "markdown", "metadata": { "id": "x-2yh9TU3FBe" }, "source": [ "### Echo Agent" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "IlZWHaxAaJSA" }, "outputs": [], "source": [ "%%writefile echo_agent/__init__.py\n", "from . import agent" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "HNhidBhKaJUL" }, "outputs": [], "source": [ "%%writefile echo_agent/agent.py\n", "import asyncio\n", "import logging\n", "\n", "from typing import Any, Dict, AsyncIterator\n", "\n", "class EchoAgent:\n", " \"\"\"\n", " A simple echo agent demonstrating bidirectional WebSocket capabilities.\n", " \"\"\"\n", " def __init__(self):\n", " self.logger = logging.getLogger(__name__)\n", " self.logger.info(\"Echo Agent ready!\")\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", " 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(\"message\", \"\")\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", "# Create an instance of our echo agent\n", "root_agent = EchoAgent()\n" ] }, { "cell_type": "markdown", "metadata": { "id": "_TrWNpCc3C7p" }, "source": [ "### API Server\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "myinCFFLXtsL" }, "outputs": [], "source": [ "%%writefile main.py\n", "import os\n", "import asyncio\n", "import logging\n", "import uvicorn\n", "from fastapi import FastAPI, WebSocket\n", "\n", "from echo_agent.agent import root_agent\n", "\n", "logging.basicConfig(level=logging.INFO)\n", "logger = logging.getLogger(\"bidi_server\")\n", "app = FastAPI()\n", "\n", "@app.websocket(\"/bidi_echo\")\n", "async def bidi_handler(websocket: WebSocket):\n", " await websocket.accept()\n", "\n", " initial_data = await websocket.receive_json(mode=\"binary\")\n", " logger.info(f\"Inbound initial handshake message: {initial_data}\")\n", "\n", " queue = asyncio.Queue()\n", " # Seed the queue if the initial frame contains a starting payload\n", " if \"input\" in initial_data:\n", " queue.put_nowait(initial_data[\"input\"])\n", "\n", " async def receive_messages():\n", " while True:\n", " try:\n", " data = await websocket.receive_json(mode=\"binary\")\n", " logger.info(f\"Inbound bidi message: {data}\")\n", "\n", " queue.put_nowait(data.get(\"input\", {}))\n", " except Exception as e:\n", " logger.info(f\"WebSocket connection closed or receive error: {e}\")\n", " break\n", "\n", " async def send_messages():\n", " try:\n", " async for response in root_agent.bidi_stream_query(queue):\n", " await websocket.send_json({\"output\": response}, mode=\"binary\")\n", " except Exception as e:\n", " logger.error(f\"Error sending bidi response: {e}\")\n", " # Run WebSocket communication loops\n", " await asyncio.gather(receive_messages(), send_messages())\n", "\n", "# You can add more FastAPI routes or configurations below if needed\n", "# Example:\n", "# @app.get(\"/hello\")\n", "# async def read_root():\n", "# return {\"Hello\": \"World\"}\n", "\n", "if __name__ == \"__main__\":\n", " # Use the PORT environment variable provided by Cloud Run, defaulting to 8080\n", " uvicorn.run(app, host=\"0.0.0.0\", port=int(os.environ.get(\"PORT\", 8080)))" ] }, { "cell_type": "markdown", "metadata": { "id": "BVuzYhGy60-b" }, "source": [ "## Deploy\n", "\n", "\n", "We make the API call by following https://docs.cloud.google.com/gemini-enterprise-agent-platform/scale/runtime/deploy-an-agent#from-dockerfile" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "sY1lWbjJajb3" }, "outputs": [], "source": [ "remote_echo_agent = client.agent_engines.create(\n", " config={\n", " \"display_name\": \"Bidi Echo test agent\",\n", " \"description\": \"bidi stream testing with echo agent\",\n", " \"source_packages\": [\n", " # The files in the current directory to upload. You can also use \".\"\n", " \"echo_agent\",\n", " \"Dockerfile\",\n", " \"main.py\",\n", " \"requirements.txt\",\n", " ],\n", " \"image_spec\": {}, # tells AgentRuntime to use the Dockerfile\n", " \"agent_framework\": \"google-adk\", # For usage through the console / UI\n", " \"env_vars\": {\"GOOGLE_GENAI_USE_VERTEXAI\": \"1\"},\n", " \"resource_limits\": {\"cpu\": \"2\", \"memory\": \"8Gi\"},\n", " # Explicitly set max_instances to stay within quota\n", " \"max_instances\": 5,\n", " # You can also adjust min_instances if needed, default is 1\n", " # \"min_instances\": 1,\n", " # # Cloud Storage bucket for staging\n", " # \"staging_bucket\": BUCKET_URI,\n", " },\n", ")\n", "remote_echo_agent" ] }, { "cell_type": "markdown", "metadata": { "id": "2GaSC4aWYsgn" }, "source": [ "## Query\n", "\n", "With a custom WebSocket server running in the deployed Agent, we will connect to the Agent Runtime API using a standard WebSocket client library rather than using the Agent Platform SDK. An example of this is shown below, sending a stream of messages to the agent before disconnecting." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "DA2qx-HeZsDW" }, "outputs": [], "source": [ "# Put the echo agent's reasoning engine id from the last step here\n", "ECHO_ENGINE_ID = \"1266316887558455296\" # @param {type: \"string\"}" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "QBU6EQbvYt7x" }, "outputs": [], "source": [ "import asyncio\n", "import json\n", "\n", "import google.auth\n", "import websockets\n", "from google.auth.transport.requests import Request\n", "\n", "creds, _ = google.auth.default()\n", "if not creds.valid:\n", " creds.refresh(Request())\n", "\n", "AGENT_PATH = \"bidi_echo\" # This should match the API server's handler path\n", "url = f\"wss://{ENDPOINT.removeprefix('https://')}/reasoningEngines/ws/internal/projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{ECHO_ENGINE_ID}/api/{AGENT_PATH}\"\n", "print(url)\n", "\n", "\n", "async def test_websocket():\n", " headers = {\n", " \"Authorization\": f\"Bearer {creds.token}\",\n", " \"Content-Type\": \"application/octet-stream\",\n", " }\n", "\n", " messages = [\"first request\", \"second request\", \"almost done!\", \"exit\"]\n", "\n", " try:\n", " async with websockets.connect(url, additional_headers=headers) as websocket:\n", " print(f\"Connected to {url}\\n\")\n", "\n", " for msg in messages:\n", " request_dict = {\"input\": {\"message\": msg}}\n", " print(f\"Sending: {request_dict}\")\n", " await websocket.send(json.dumps(request_dict))\n", "\n", " # Listen for response\n", " try:\n", " response = await asyncio.wait_for(websocket.recv(), timeout=10.0)\n", " print(f\"Received: {response}\\n\")\n", " except asyncio.TimeoutError:\n", " print(\"No message received within the timeout period.\\n\")\n", "\n", " except Exception as e:\n", " print(f\"An error occurred: {e}\")\n", " raise e\n", "\n", "\n", "await test_websocket()" ] }, { "cell_type": "markdown", "metadata": { "id": "H6x4TBsqbGr_" }, "source": [ "# Live API Agent\n", "\n", "We now show how to use the bidirectional streaming API with a Live ADK Agent for real-time audio streaming." ] }, { "cell_type": "markdown", "metadata": { "id": "3G5hcLXmbeLP" }, "source": [ "## Setting up the source files\n", "\n", "Once again, we will deploy using Agent Runtime's bring-your-own-Docker image deployment.\n", "\n", "We define the following files and subdirectories in the current working directory.\n", "\n", "* `capital_agent/`: the AI logic for the Capital Agent. It contains the following files:\n", " * `__init__.py`\n", " * `agent.py`\n", "* `main.py`: the API server to run\n", "* `requirements.txt`: the PyPI dependencies to pick up\n", "* `Dockerfile`: the commands to assemble the image" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "XtS21V96b8p9" }, "outputs": [], "source": [ "!mkdir -p capital_agent" ] }, { "cell_type": "markdown", "metadata": { "id": "XglhKKHWb3Je" }, "source": [ "### Requirements" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "2Xv0hndmb4wp" }, "outputs": [], "source": [ "%%writefile requirements.txt\n", "google-cloud-aiplatform[agent_engines,adk]\n", "fastapi\n", "uvicorn\n", "pydantic" ] }, { "cell_type": "markdown", "metadata": { "id": "DQiC3imncLqN" }, "source": [ "### Dockerfile" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tvAC9Ma9cMxi" }, "outputs": [], "source": [ "%%writefile Dockerfile\n", "FROM python:3.13-alpine\n", "WORKDIR /app\n", "\n", "RUN adduser --disabled-password --gecos \"\" myuser\n", "\n", "COPY requirements.txt .\n", "RUN pip install --no-cache-dir -r requirements.txt\n", "\n", "\n", "COPY --chown=myuser:myuser . .\n", "\n", "ENV PATH=\"/home/myuser/.local/bin:$PATH\"\n", "\n", "USER myuser\n", "\n", "CMD [\"sh\", \"-c\", \"uvicorn main:app --host 0.0.0.0 --port $PORT\"]" ] }, { "cell_type": "markdown", "metadata": { "id": "CkEBQ9x3cQgx" }, "source": [ "### Capital Agent" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "-kDoOWKNcRd0" }, "outputs": [], "source": [ "%%writefile capital_agent/__init__.py\n", "from . import agent" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "fpyXBf0GcVtn" }, "outputs": [], "source": [ "%%writefile capital_agent/agent.py\n", "from google.adk.agents import Agent\n", "\n", "# Define a tool function\n", "def get_capital_city(country: str) -> str:\n", " \"\"\"Retrieves the capital city for a given country.\"\"\"\n", " # Replace with actual logic (e.g., API call, database lookup)\n", " capitals = {\"france\": \"Paris\", \"japan\": \"Tokyo\", \"canada\": \"Ottawa\"}\n", " return capitals.get(country.lower(), f\"Sorry, I don't know the capital of {country}.\")\n", "\n", "root_agent = Agent(\n", " model=\"gemini-live-2.5-flash-native-audio\", # Specify the Live API model\n", " name=\"capital_agent\",\n", " description=\"Answers user questions about the capital city of a given country using Live API.\",\n", " instruction=\"\"\"You are an agent that provides the capital city of a country.\n", " When a user asks for the capital of a country:\n", " 1. Identify the country name from the user's query.\n", " 2. Use the `get_capital_city` tool to find the capital.\n", " 3. Respond clearly to the user, stating the capital city.\n", " Example Query: \"What's the capital of France?\"\n", " Example Response: \"The capital of France is Paris.\"\n", " \"\"\",\n", " tools=[get_capital_city]\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "13yehF0QcfZw" }, "source": [ "### API Server" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "KB60oo6LchQU" }, "outputs": [], "source": [ "%%writefile main.py\n", "import os\n", "\n", "import uvicorn\n", "from fastapi import FastAPI\n", "from google.adk.cli.fast_api import get_fast_api_app\n", "\n", "# Get the directory where main.py is located\n", "AGENT_DIR = os.path.dirname(os.path.abspath(__file__))\n", "# Example session service URI (e.g., SQLite)\n", "SESSION_SERVICE_URI = \"\"\n", "# Example allowed origins for CORS\n", "ALLOWED_ORIGINS = [\"http://localhost\", \"http://localhost:8080\", \"*\"]\n", "# Set web=True if you intend to serve a web interface, False otherwise\n", "SERVE_WEB_INTERFACE = True\n", "\n", "# Call the function to get the FastAPI app instance\n", "# Ensure the agent directory name ('capital_agent') matches your agent folder\n", "app: FastAPI = get_fast_api_app(\n", " agents_dir=AGENT_DIR,\n", " session_service_uri=SESSION_SERVICE_URI,\n", " allow_origins=ALLOWED_ORIGINS,\n", " web=SERVE_WEB_INTERFACE,\n", ")\n", "\n", "# You can add more FastAPI routes or configurations below if needed\n", "# Example:\n", "# @app.get(\"/hello\")\n", "# async def read_root():\n", "# return {\"Hello\": \"World\"}\n", "\n", "if __name__ == \"__main__\":\n", " # Use the PORT environment variable provided by Cloud Run, defaulting to 8080\n", " uvicorn.run(app, host=\"0.0.0.0\", port=int(os.environ.get(\"PORT\", 8080)))" ] }, { "cell_type": "markdown", "metadata": { "id": "EZV8BuKZcq0h" }, "source": [ "## Deploy\n", "\n", "We make the API call by following https://docs.cloud.google.com/gemini-enterprise-agent-platform/scale/runtime/deploy-an-agent#from-dockerfile" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "OllI3G8Vcs2-" }, "outputs": [], "source": [ "remote_live_agent = client.agent_engines.create(\n", " config={\n", " \"display_name\": \"Bidi Live API test agent\",\n", " \"description\": \"bidi stream testing with live agent\",\n", " \"source_packages\": [\n", " # The files in the current directory to upload.\n", " \"capital_agent\",\n", " \"Dockerfile\",\n", " \"main.py\",\n", " \"requirements.txt\",\n", " ],\n", " \"image_spec\": {}, # tells AgentRuntime to use the Dockerfile\n", " \"agent_framework\": \"google-adk\", # For usage through the console / UI\n", " \"env_vars\": {\"GOOGLE_GENAI_USE_VERTEXAI\": \"1\"},\n", " \"resource_limits\": {\"cpu\": \"2\", \"memory\": \"8Gi\"},\n", " # Explicitly set max_instances to stay within quota\n", " \"max_instances\": 5,\n", " # You can also adjust min_instances if needed, default is 1\n", " # \"min_instances\": 1,\n", " },\n", ")\n", "remote_live_agent" ] }, { "cell_type": "markdown", "metadata": { "id": "eM9qYzCOdTSm" }, "source": [ "## Connect to the Agent" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Ay9Siu_MetgV" }, "outputs": [], "source": [ "# Put the live agent's reasoning engine id from the last step here\n", "LIVE_ENGINE_ID = \"7678296516062085120\" # @param {type: \"string\"}" ] }, { "cell_type": "markdown", "metadata": { "id": "aqefAU0of1h1" }, "source": [ "### Create Session\n", "\n", "For the Live API Agent, you first need to create a session. For our example agent, the session service is running in-memory on the agent itself. We can connect via HTTP to the built-in `/apps/capital_agent/users/user-123/sessions` endpoint to create the session." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "uWN48DHYfVvs" }, "outputs": [], "source": [ "APP_NAME = \"capital_agent\"\n", "USER_ID = \"demo_user\"\n", "\n", "# Put the session id you want to create\n", "SESSION_ID = \"demo_session\" # @param {type: \"string\"}" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "WaILod34qXoN" }, "outputs": [], "source": [ "import google.auth\n", "import requests\n", "from google.auth.transport.requests import Request\n", "\n", "creds, _ = google.auth.default()\n", "if not creds.valid:\n", " creds.refresh(Request())\n", "\n", "base_ingress_url = f\"{ENDPOINT}/reasoningEngines/internal/projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{LIVE_ENGINE_ID}/api\"\n", "create_session_url = (\n", " f\"{base_ingress_url}/apps/{APP_NAME}/users/{USER_ID}/sessions/{SESSION_ID}\"\n", ")\n", "\n", "headers = {\"Authorization\": f\"Bearer {creds.token}\", \"Content-Type\": \"application/json\"}\n", "\n", "response = requests.post(create_session_url, headers=headers)\n", "response.json()" ] }, { "cell_type": "markdown", "metadata": { "id": "c9S5YE3jtoRZ" }, "source": [ "### List Sessions" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ikfkZ3NHtp44" }, "outputs": [], "source": [ "import google.auth\n", "import requests\n", "from google.auth.transport.requests import Request\n", "\n", "creds, _ = google.auth.default()\n", "if not creds.valid:\n", " creds.refresh(Request())\n", "\n", "base_ingress_url = f\"{ENDPOINT}/reasoningEngines/internal/projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{LIVE_ENGINE_ID}/api\"\n", "list_sessions_url = f\"{base_ingress_url}/apps/{APP_NAME}/users/{USER_ID}/sessions/\"\n", "\n", "headers = {\"Authorization\": f\"Bearer {creds.token}\", \"Content-Type\": \"application/json\"}\n", "\n", "response = requests.get(list_sessions_url, headers=headers)\n", "response.json()" ] }, { "cell_type": "markdown", "metadata": { "id": "yXL77455gE6m" }, "source": [ "### Query\n", "\n", "This example sends a request to the agent and receives an audio response over the bidirectional WebSocket connection." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "QdEHaxCkdvxA" }, "outputs": [], "source": [ "LIVE_QUERY_TEXT = \"What is the capital of France?\" # @param {type: \"string\"}" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "sMH2TL5Ddr5l" }, "outputs": [], "source": [ "import base64\n", "import time\n", "import wave\n", "\n", "import google.auth\n", "from IPython.display import Audio, display\n", "from google.auth.transport.requests import Request\n", "\n", "creds, _ = google.auth.default()\n", "if not creds.valid:\n", " creds.refresh(Request())\n", "\n", "# Parameters for the audio\n", "sample_rate = 44100 # samples per second\n", "\n", "url = f\"wss://{ENDPOINT.removeprefix('https://')}/reasoningEngines/ws/internal/projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{LIVE_ENGINE_ID}/api/run_live?app_name={APP_NAME}&user_id={USER_ID}&session_id={SESSION_ID}\"\n", "\n", "\n", "def save_wav_file(filename, pcm_data, channels=1, sampwidth=2, framerate=24000):\n", " \"\"\"Saves raw PCM data to a WAV file.\"\"\"\n", " with wave.open(filename, \"wb\") as wf:\n", " wf.setnchannels(channels)\n", " wf.setsampwidth(sampwidth) # 2 bytes for 16-bit audio\n", " wf.setframerate(framerate)\n", " wf.writeframes(pcm_data)\n", " print(f\"Audio saved to {filename}\")\n", "\n", "\n", "async def test_websocket():\n", " headers = {\n", " \"Authorization\": f\"Bearer {creds.token}\",\n", " }\n", "\n", " try:\n", " async with websockets.connect(url, additional_headers=headers) as websocket:\n", " print(f\"Connected to {url}\")\n", " test_message = {\n", " \"content\": {\"role\": \"user\", \"parts\": [{\"text\": LIVE_QUERY_TEXT}]}\n", " }\n", " print(f\"Sending: {test_message}\")\n", " await websocket.send(json.dumps(test_message))\n", " print(f\"Sent: {test_message}\")\n", "\n", " # Listen for responses\n", " print(\"--- Receiving Responses ---\")\n", " audio_buffer = b\"\"\n", " text_response = \"\"\n", " try:\n", " while True:\n", " response = await asyncio.wait_for(websocket.recv(), timeout=10.0)\n", " print(f\"Received: {response}\")\n", " try:\n", " data = json.loads(response)\n", "\n", " if \"content\" in data and data[\"content\"][\"role\"] == \"model\":\n", " for part in data[\"content\"].get(\"parts\", []):\n", " if \"text\" in part:\n", " text_response += part[\"text\"]\n", " print(\n", " f\"====== MODEL TEXT PART: {part['text']} ======\"\n", " )\n", " if \"inlineData\" in part:\n", " mime_type = part[\"inlineData\"].get(\"mimeType\")\n", " if mime_type == \"audio/pcm\":\n", " b64_data = part[\"inlineData\"].get(\"data\")\n", " if b64_data:\n", " try:\n", " print(type(b64_data))\n", " decoded_chunk = (\n", " base64.urlsafe_b64decode(b64_data)\n", " )\n", " audio_buffer += decoded_chunk\n", " print(\n", " f\"Decoded and appended {len(decoded_chunk)} audio bytes.\"\n", " )\n", " except base64.binascii.Error as e:\n", " print(b64_data)\n", " print(\n", " f\"Base64 Decode Error: {e} - on chunk length {len(b64_data)}\"\n", " )\n", "\n", " if data.get(\"turnComplete\"):\n", " print(\"Turn complete received.\")\n", " if text_response:\n", " print(\n", " f\"====== FINAL TEXT RESPONSE: {text_response} ======\"\n", " )\n", " if audio_buffer:\n", " print(\"Saving audio buffer\")\n", " filename = f\"response_{int(time.time())}.wav\"\n", " save_wav_file(filename, audio_buffer)\n", " display(Audio(filename=filename, rate=sample_rate))\n", " audio_buffer = b\"\"\n", " break\n", " except json.JSONDecodeError:\n", " print(f\"Received non-JSON message: {response}\")\n", " except asyncio.TimeoutError:\n", " print(\"No message received within the timeout period.\")\n", " except websockets.exceptions.ConnectionClosedOK:\n", " print(\"Connection closed normally.\")\n", " except websockets.exceptions.ConnectionClosedError as e:\n", " print(f\"Connection closed with error: {e}\")\n", "\n", " except Exception as e:\n", " print(f\"An error occurred: {e}\")\n", " raise e\n", "\n", "\n", "await test_websocket()" ] }, { "cell_type": "markdown", "metadata": { "id": "MXJ1DSBg23g5" }, "source": [ "# Clean up" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "M3d6JGUa6Mp7" }, "outputs": [], "source": [ "remote_echo_agent.delete(force=True)\n", "remote_live_agent.delete(force=True)" ] } ], "metadata": { "colab": { "name": "tutorial_bidi_stream_v2.ipynb", "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }