wehub-resource-sync
1118 lines
40 KiB
Plaintext
1118 lines
40 KiB
Plaintext
{
|
|
"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",
|
|
"<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_bidi_stream_v2.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/agent-platform/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fagents%2Fagent_engine%2Ftutorial_bidi_stream_v2.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/agent-platform/workbench/instances?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/agents/agent_engine/tutorial_bidi_stream_v2.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://storage.googleapis.com/github-repo/workbench-icon.svg\" alt=\"Workbench logo\"><br> Open in 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_bidi_stream_v2.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",
|
|
"<p>\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_bidi_stream_v2.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_bidi_stream_v2.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_bidi_stream_v2.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_bidi_stream_v2.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_bidi_stream_v2.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>\n",
|
|
"</p>"
|
|
]
|
|
},
|
|
{
|
|
"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
|
|
}
|