agent-starter-pack 0.15.6__py3-none-any.whl → 0.16.0__py3-none-any.whl

src/deployment_targets/agent_engine/tests/load_test/load_test.py

@@ -11,6 +11,134 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+{%- if cookiecutter.agent_name == "adk_live" %}
+
+import json
+import logging
+import time
+
+from locust import User, between, task
+from websockets.exceptions import WebSocketException
+from websockets.sync.client import connect
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class WebSocketUser(User):
+    """Simulates a user making websocket requests to the remote agent engine."""
+
+    wait_time = between(1, 3)  # Wait 1-3 seconds between tasks
+    abstract = True
+
+    def __init__(self, *args: object, **kwargs: object) -> None:
+        super().__init__(*args, **kwargs)
+        if self.host.startswith("https://"):
+            self.ws_url = self.host.replace("https://", "wss://", 1) + "/ws"
+        elif self.host.startswith("http://"):
+            self.ws_url = self.host.replace("http://", "ws://", 1) + "/ws"
+        else:
+            self.ws_url = self.host + "/ws"
+
+    @task
+    def websocket_audio_conversation(self) -> None:
+        """Test a full websocket conversation with audio input."""
+        start_time = time.time()
+        response_count = 0
+        exception = None
+
+        try:
+            response_count = self._websocket_interaction()
+
+            # Mark as failure if we got no valid responses
+            if response_count == 0:
+                exception = Exception("No responses received from agent")
+
+        except WebSocketException as e:
+            exception = e
+            logger.error(f"WebSocket error: {e}")
+        except Exception as e:
+            exception = e
+            logger.error(f"Unexpected error: {e}")
+        finally:
+            total_time = int((time.time() - start_time) * 1000)
+
+            # Report the request metrics to Locust
+            self.environment.events.request.fire(
+                request_type="WS",
+                name="websocket_conversation",
+                response_time=total_time,
+                response_length=response_count * 100,  # Approximate response size
+                response=None,
+                context={},
+                exception=exception,
+            )
+
+    def _websocket_interaction(self) -> int:
+        """Handle the websocket interaction and return response count."""
+        response_count = 0
+
+        with connect(self.ws_url, open_timeout=10, close_timeout=20) as websocket:
+            # Wait for setupComplete
+            setup_response = websocket.recv(timeout=10.0)
+            setup_data = json.loads(setup_response)
+            assert "setupComplete" in setup_data, (
+                f"Expected setupComplete, got {setup_data}"
+            )
+            logger.info("Received setupComplete")
+
+            # Send dummy audio chunk with user_id
+            dummy_audio = bytes([0] * 1024)
+            audio_msg = {
+                "user_id": "load-test-user",
+                "realtimeInput": {
+                    "mediaChunks": [
+                        {
+                            "mimeType": "audio/pcm;rate=16000",
+                            "data": dummy_audio.hex(),
+                        }
+                    ]
+                },
+            }
+            websocket.send(json.dumps(audio_msg))
+            logger.info("Sent audio chunk")
+
+            # Send text message to complete the turn
+            text_msg = {
+                "content": {
+                    "role": "user",
+                    "parts": [{"text": "Hello!"}],
+                }
+            }
+            websocket.send(json.dumps(text_msg))
+            logger.info("Sent text completion")
+
+            # Collect responses until turn_complete or timeout
+            for _ in range(20):  # Max 20 responses
+                try:
+                    response = websocket.recv(timeout=10.0)
+                    response_data = json.loads(response)
+                    response_count += 1
+                    logger.debug(f"Received response: {response_data}")
+
+                    if isinstance(response_data, dict) and response_data.get(
+                        "turn_complete"
+                    ):
+                        logger.info(f"Turn complete after {response_count} responses")
+                        break
+                except TimeoutError:
+                    logger.info(f"Timeout after {response_count} responses")
+                    break
+
+        return response_count
+
+
+class RemoteAgentUser(WebSocketUser):
+    """User for testing remote agent engine deployment."""
+
+    # Set the host via command line: locust -f load_test.py --host=https://your-deployed-service.run.app
+    host = "http://localhost:8000"  # Default for local testing
+{%- else %}
 
 import json
 import logging
@@ -54,7 +182,7 @@ class ChatStreamUser(HttpUser):
         """Simulates a chat stream interaction."""
         headers = {"Content-Type": "application/json"}
         headers["Authorization"] = f"Bearer {os.environ['_AUTH_TOKEN']}"
-{% if "adk" in cookiecutter.tags %}
+{% if cookiecutter.is_adk %}
         data = {
             "class_method": "async_stream_query",
             "input": {
@@ -84,7 +212,7 @@ class ChatStreamUser(HttpUser):
             headers=headers,
             json=data,
             catch_response=True,
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
             name="/streamQuery async_stream_query",
 {%- else %}
             name="/stream_messages first message",
@@ -112,7 +240,7 @@ class ChatStreamUser(HttpUser):
                 total_time = end_time - start_time
                 self.environment.events.request.fire(
                     request_type="POST",
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
                     name="/streamQuery end",
 {%- else %}
                     name="/stream_messages end",
@@ -124,3 +252,4 @@ class ChatStreamUser(HttpUser):
                 )
             else:
                 response.failure(f"Unexpected status code: {response.status_code}")
+{%- endif %}

src/deployment_targets/agent_engine/{{cookiecutter.agent_directory}}/agent_engine_app.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # mypy: disable-error-code="attr-defined,arg-type"
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
 import logging
 import os
 from typing import Any
@@ -25,8 +25,12 @@ from google.adk.artifacts import GcsArtifactService
 from google.cloud import logging as google_cloud_logging
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider, export
-from vertexai._genai.types import AgentEngine, AgentEngineConfig
+from vertexai._genai.types import AgentEngine, AgentEngineConfig{%- if cookiecutter.is_adk_live %}, AgentServerMode{%- endif %}
+{%- if cookiecutter.is_adk_live %}
+from vertexai.preview.reasoning_engines import AdkApp
+{%- else %}
 from vertexai.agent_engines.templates.adk import AdkApp
+{%- endif %}
 
 from {{cookiecutter.agent_directory}}.agent import root_agent
 from {{cookiecutter.agent_directory}}.utils.deployment import (
@@ -282,7 +286,7 @@ def deploy_agent_engine_app(
     # Read requirements
     with open(requirements_file) as f:
         requirements = f.read().strip().split("\n")
-{% if "adk" in cookiecutter.tags %}
+{% if cookiecutter.is_adk %}
     agent_engine = AgentEngineApp(
         agent=root_agent,
         artifact_service_builder=lambda: GcsArtifactService(
@@ -314,6 +318,10 @@ def deploy_agent_engine_app(
         requirements=requirements,
         staging_bucket=staging_bucket_uri,
         labels=labels,
+{%- if cookiecutter.is_adk_live %}
+        agent_server_mode=AgentServerMode.EXPERIMENTAL,  # Enable bidi streaming
+        resource_limits={"cpu": "4", "memory": "8Gi"},
+{%- endif %}
     )
 
     agent_config = {

src/deployment_targets/agent_engine/{{cookiecutter.agent_directory}}/utils/deployment.py

@@ -75,9 +75,13 @@ def print_deployment_success(
     agent_engine_id = remote_agent.api_resource.name.split("/")[-1]
     console_url = f"https://console.cloud.google.com/vertex-ai/agents/locations/{location}/agent-engines/{agent_engine_id}?project={project}"
 
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
     print(
+{%- if cookiecutter.is_adk_live %}
+        f"\n✅ Deployment successful! Run your agent with: `make playground-remote`"
+{%- else %}
         f"\n✅ Deployment successful! Test your agent: notebooks/adk_app_testing.ipynb"
+{%- endif %}
         f"\n📊 View in console: {console_url}\n"
     )
 {%- else %}

src/deployment_targets/agent_engine/{{cookiecutter.agent_directory}}/utils/{% if cookiecutter.is_adk_live %}expose_app.py{% else %}unused_expose_app.py{% endif %}

@@ -0,0 +1,461 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import json
+import logging
+from collections.abc import Callable
+from pathlib import Path
+from typing import Literal
+
+import backoff
+import google.auth
+import vertexai
+from fastapi import FastAPI, HTTPException, WebSocket
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+from google.cloud import logging as google_cloud_logging
+from pydantic import BaseModel
+from websockets.exceptions import ConnectionClosedError
+
+from ..agent import root_agent
+from ..agent_engine_app import AgentEngineApp
+
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Get the path to the frontend build directory
+current_dir = Path(__file__).parent
+frontend_build_dir = current_dir.parent.parent / "frontend" / "build"
+
+# Mount static files if build directory exists
+if frontend_build_dir.exists():
+    app.mount(
+        "/static",
+        StaticFiles(directory=str(frontend_build_dir / "static")),
+        name="static",
+    )
+logging_client = google_cloud_logging.Client()
+logger = logging_client.logger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# Global configuration for agent engine mode
+USE_REMOTE_AGENT = False
+REMOTE_AGENT_ENGINE_ID = None
+PROJECT_ID = None
+LOCATION = "us-central1"
+LOCAL_AGENT_CALLABLE = None
+
+
+class WebSocketToQueueAdapter:
+    """Adapter to convert WebSocket messages to an asyncio Queue for the agent engine."""
+
+    def __init__(
+        self, websocket: WebSocket, agent_engine: AgentEngineApp | None = None
+    ):
+        """Initialize the adapter.
+
+        Args:
+            websocket: The client websocket connection
+            agent_engine: The agent engine instance (None if using remote)
+        """
+        self.websocket = websocket
+        self.agent_engine = agent_engine
+        self.input_queue: asyncio.Queue[dict] = asyncio.Queue()
+        self.first_message = True
+
+    def _transform_remote_agent_engine_response(self, response: dict) -> dict:
+        """Transform remote Agent Engine bidiStreamOutput to ADK Event format for frontend."""
+        # Check if this is a remote Agent Engine bidiStreamOutput
+        bidi_output = response.get("bidiStreamOutput")
+        if not bidi_output:
+            # Not a remote agent engine response, return as-is
+            return response
+
+        # Transform to ADK Event format that frontend already handles
+        # Just unwrap the bidiStreamOutput wrapper - the content is already in ADK Event format
+        return bidi_output
+
+    async def receive_from_client(self) -> None:
+        """Listen for messages from the client and put them in the queue."""
+        while True:
+            try:
+                # Use receive() instead of receive_json() to handle both text and binary data
+                message = await self.websocket.receive()
+
+                # Handle different message types
+                if "text" in message:
+                    # Parse JSON text messages
+                    data = json.loads(message["text"])
+
+                    if isinstance(data, dict):
+                        # Skip setup messages - they're for backend logging only, not valid LiveRequest format
+                        if "setup" in data:
+                            # Log setup information
+                            logger.log_struct(
+                                {**data["setup"], "type": "setup"}, severity="INFO"
+                            )
+                            logging.info(
+                                "Received setup message (not forwarding to agent)"
+                            )
+                            continue
+
+                        # Frontend handles message format for both modes
+                        await self.input_queue.put(data)
+                    else:
+                        logging.warning(
+                            f"Received unexpected JSON structure from client: {data}"
+                        )
+
+                elif "bytes" in message:
+                    # Handle binary data
+                    # Convert binary to appropriate format for agent engine
+                    await self.input_queue.put({"binary_data": message["bytes"]})
+
+                else:
+                    logging.warning(
+                        f"Received unexpected message type from client: {message}"
+                    )
+
+            except ConnectionClosedError as e:
+                logging.warning(f"Client closed connection: {e}")
+                break
+            except json.JSONDecodeError as e:
+                logging.error(f"Error parsing JSON from client: {e}")
+                break
+            except Exception as e:
+                logging.error(f"Error receiving from client: {e!s}")
+                break
+
+    async def run_agent_engine(self) -> None:
+        """Run the agent engine with the input queue."""
+        try:
+            if self.agent_engine is not None:
+                # Local agent engine mode
+                # Give the agent engine a moment to initialize before sending setupComplete
+                await asyncio.sleep(1)
+
+                # Send setupComplete after initialization delay
+                setup_complete_response: dict = {"setupComplete": {}}
+                await self.websocket.send_json(setup_complete_response)
+
+                async for response in self.agent_engine.bidi_stream_query(
+                    self.input_queue
+                ):
+                    # Send responses from agent engine to the websocket client
+                    if response is not None:
+                        await self.websocket.send_json(response)
+
+                        # Check for error responses
+                        if isinstance(response, dict) and "error" in response:
+                            logging.error(f"Agent engine error: {response['error']}")
+                            break
+            else:
+                # Remote agent engine mode
+                # Don't send setupComplete until remote connection is established
+                await self.run_remote_agent_engine()
+        except Exception as e:
+            logging.error(f"Error in agent engine: {e}")
+            await self.websocket.send_json({"error": str(e)})
+
+    async def run_remote_agent_engine(self) -> None:
+        """Run the remote agent engine connection."""
+        client = vertexai.Client(
+            project=PROJECT_ID,
+            location=LOCATION,
+        )
+
+        async with client.aio.live.agent_engines.connect(
+            agent_engine=REMOTE_AGENT_ENGINE_ID,
+            config={"class_method": "bidi_stream_query"},
+        ) as session:
+            # Send setupComplete only after remote connection is established
+            logging.info("Remote agent engine connection established")
+            setup_complete_response: dict = {"setupComplete": {}}
+            await self.websocket.send_json(setup_complete_response)
+
+            # Create task to forward messages from queue to remote session
+            async def forward_to_remote() -> None:
+                while True:
+                    try:
+                        message = await self.input_queue.get()
+                        await session.send(message)
+                    except Exception as e:
+                        logging.error(f"Error forwarding to remote: {e}")
+                        break
+
+            # Create task to receive from remote and send to websocket
+            async def receive_from_remote() -> None:
+                while True:
+                    try:
+                        response = await session.receive()
+                        if response is not None:
+                            # Transform remote Agent Engine bidiStreamOutput format to frontend format
+                            transformed = self._transform_remote_agent_engine_response(
+                                response
+                            )
+                            if transformed:
+                                await self.websocket.send_json(transformed)
+
+                            # Check for error responses
+                            if isinstance(response, dict) and "error" in response:
+                                logging.error(
+                                    f"Remote agent engine error: {response['error']}"
+                                )
+                                break
+                    except Exception as e:
+                        logging.error(f"Error receiving from remote: {e}")
+                        break
+
+            await asyncio.gather(
+                forward_to_remote(),
+                receive_from_remote(),
+            )
+
+
+def get_connect_and_run_callable(websocket: WebSocket) -> Callable:
+    """Create a callable that handles agent engine connection with retry logic.
+
+    Args:
+        websocket: The client websocket connection
+
+    Returns:
+        Callable: An async function that establishes and manages the agent engine connection
+    """
+
+    async def on_backoff(details: backoff._typing.Details) -> None:
+        await websocket.send_json(
+            {
+                "status": f"Model connection error, retrying in {details['wait']} seconds..."
+            }
+        )
+
+    @backoff.on_exception(
+        backoff.expo, ConnectionClosedError, max_tries=10, on_backoff=on_backoff
+    )
+    async def connect_and_run() -> None:
+        if USE_REMOTE_AGENT:
+            # Remote agent engine mode
+            logging.info(f"Connecting to remote agent engine: {REMOTE_AGENT_ENGINE_ID}")
+            adapter = WebSocketToQueueAdapter(websocket, agent_engine=None)
+        else:
+            # Local agent engine mode
+            agent_callable = (
+                LOCAL_AGENT_CALLABLE if LOCAL_AGENT_CALLABLE is not None else root_agent
+            )
+            logging.info(
+                f"Starting local agent engine with callable: {agent_callable.__name__ if hasattr(agent_callable, '__name__') else agent_callable}"
+            )
+            agent_engine = AgentEngineApp(agent=agent_callable)
+            adapter = WebSocketToQueueAdapter(websocket, agent_engine)
+
+        logging.info("Starting bidirectional communication with agent engine")
+        await asyncio.gather(
+            adapter.receive_from_client(),
+            adapter.run_agent_engine(),
+        )
+
+    return connect_and_run
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket) -> None:
+    """Handle new websocket connections."""
+    await websocket.accept()
+    connect_and_run = get_connect_and_run_callable(websocket)
+    await connect_and_run()
+
+
+class Feedback(BaseModel):
+    """Represents feedback for a conversation."""
+
+    score: int | float
+    text: str | None = ""
+    run_id: str
+    user_id: str | None
+    log_type: Literal["feedback"] = "feedback"
+
+
+@app.post("/feedback")
+def collect_feedback(feedback: Feedback) -> dict[str, str]:
+    """Collect and log feedback.
+
+    Args:
+        feedback: The feedback data to log
+
+    Returns:
+        Success message
+    """
+    logger.log_struct(feedback.model_dump(), severity="INFO")
+    return {"status": "success"}
+
+
+@app.get("/")
+async def serve_frontend_root() -> FileResponse:
+    """Serve the frontend index.html at the root path."""
+    index_file = frontend_build_dir / "index.html"
+    if index_file.exists():
+        return FileResponse(str(index_file))
+    raise HTTPException(
+        status_code=404,
+        detail="Frontend not built. Run 'npm run build' in the frontend directory.",
+    )
+
+
+@app.get("/{full_path:path}")
+async def serve_frontend_spa(full_path: str) -> FileResponse:
+    """Catch-all route to serve the frontend for SPA routing.
+
+    This ensures that client-side routes are handled by the React app.
+    Excludes API routes (ws, feedback) and static assets.
+    """
+    # Don't intercept API routes
+    if full_path.startswith(("ws", "feedback", "static", "api")):
+        raise HTTPException(status_code=404, detail="Not found")
+
+    # Serve index.html for all other routes (SPA routing)
+    index_file = frontend_build_dir / "index.html"
+    if index_file.exists():
+        return FileResponse(str(index_file))
+    raise HTTPException(
+        status_code=404,
+        detail="Frontend not built. Run 'npm run build' in the frontend directory.",
+    )
+
+
+# Main execution
+if __name__ == "__main__":
+    import argparse
+
+    import uvicorn
+
+    parser = argparse.ArgumentParser(description="Agent Engine Proxy Server")
+    parser.add_argument(
+        "--mode",
+        choices=["local", "remote"],
+        default="local",
+        help="Agent engine mode: 'local' for local agent or 'remote' for deployed agent engine",
+    )
+    parser.add_argument(
+        "--remote-id",
+        type=str,
+        help="Remote agent engine ID (required when mode=remote)",
+    )
+    parser.add_argument(
+        "--project-id", type=str, help="GCP project ID (required when mode=remote)"
+    )
+    parser.add_argument(
+        "--location",
+        type=str,
+        default="us-central1",
+        help="GCP location (default: us-central1)",
+    )
+    parser.add_argument(
+        "--local-agent",
+        type=str,
+        help="Python path to local agent callable (e.g., 'app.agent.root_agent'). Defaults to root_agent",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="Port to run the server on (default: 8000)",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="localhost",
+        help="Host to run the server on (default: localhost)",
+    )
+
+    args = parser.parse_args()
+
+    if args.mode == "remote":
+        USE_REMOTE_AGENT = True
+
+        # Try to load from deployment_metadata.json if remote-id not provided
+        if not args.remote_id:
+            deployment_metadata_path = (
+                Path(__file__).parent.parent.parent / "deployment_metadata.json"
+            )
+            if deployment_metadata_path.exists():
+                with open(deployment_metadata_path) as f:
+                    metadata = json.load(f)
+                    REMOTE_AGENT_ENGINE_ID = metadata.get("remote_agent_engine_id")
+                    if not REMOTE_AGENT_ENGINE_ID:
+                        parser.error(
+                            "No remote_agent_engine_id found in deployment_metadata.json"
+                        )
+                    print("Loaded remote agent engine ID from deployment_metadata.json")
+            else:
+                parser.error(
+                    "--remote-id is required when deployment_metadata.json is not found"
+                )
+        else:
+            REMOTE_AGENT_ENGINE_ID = args.remote_id
+
+        # Extract project ID from remote agent engine ID if not provided
+        if not args.project_id:
+            # Format: projects/PROJECT_ID/locations/LOCATION/reasoningEngines/ENGINE_ID
+            import re
+
+            match = re.match(
+                r"projects/([^/]+)/locations/([^/]+)/reasoningEngines/",
+                REMOTE_AGENT_ENGINE_ID,
+            )
+            if match:
+                PROJECT_ID = match.group(1)
+                extracted_location = match.group(2)
+                LOCATION = (
+                    args.location
+                    if args.location != "us-central1"
+                    else extracted_location
+                )
+                print("Extracted project ID and location from remote agent engine ID")
+            else:
+                # Fall back to google.auth.default()
+                try:
+                    _, PROJECT_ID = google.auth.default()
+                    LOCATION = args.location
+                    print(f"Using default project ID from google.auth: {PROJECT_ID}")
+                except Exception as e:
+                    parser.error(f"Could not determine project ID: {e}")
+        else:
+            PROJECT_ID = args.project_id
+            LOCATION = args.location
+
+        print("Starting server in REMOTE mode:")
+        print(f"  Remote Agent Engine ID: {REMOTE_AGENT_ENGINE_ID}")
+        print(f"  Project ID: {PROJECT_ID}")
+        print(f"  Location: {LOCATION}")
+    else:
+        print("Starting server in LOCAL mode")
+        if args.local_agent:
+            # Dynamically import the agent callable
+            import importlib
+
+            module_path, callable_name = args.local_agent.rsplit(".", 1)
+            module = importlib.import_module(module_path)
+            LOCAL_AGENT_CALLABLE = getattr(module, callable_name)
+            print(f"  Using custom agent: {args.local_agent}")
+        else:
+            print("  Using default agent: root_agent")
+
+    uvicorn.run(app, host=args.host, port=args.port)