agent-starter-pack 0.15.7__py3-none-any.whl → 0.17.0__py3-none-any.whl

agent_starter_pack/deployment_targets/agent_engine/tests/load_test/load_test.py

@@ -0,0 +1,255 @@
+# 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.
+{%- 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
+import os
+import time
+
+from locust import HttpUser, between, task
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+# Initialize Vertex AI and load agent config
+with open("deployment_metadata.json") as f:
+    remote_agent_engine_id = json.load(f)["remote_agent_engine_id"]
+
+parts = remote_agent_engine_id.split("/")
+project_id = parts[1]
+location = parts[3]
+engine_id = parts[5]
+
+# Convert remote agent engine ID to streaming URL.
+base_url = f"https://{location}-aiplatform.googleapis.com"
+url_path = f"/v1/projects/{project_id}/locations/{location}/reasoningEngines/{engine_id}:streamQuery"
+
+logger.info("Using remote agent engine ID: %s", remote_agent_engine_id)
+logger.info("Using base URL: %s", base_url)
+logger.info("Using URL path: %s", url_path)
+
+
+class ChatStreamUser(HttpUser):
+    """Simulates a user interacting with the chat stream API."""
+
+    wait_time = between(1, 3)  # Wait 1-3 seconds between tasks
+    host = base_url  # Set the base host URL for Locust
+
+    @task
+    def chat_stream(self) -> None:
+        """Simulates a chat stream interaction."""
+        headers = {"Content-Type": "application/json"}
+        headers["Authorization"] = f"Bearer {os.environ['_AUTH_TOKEN']}"
+{% if cookiecutter.is_adk %}
+        data = {
+            "class_method": "async_stream_query",
+            "input": {
+                "user_id": "test",
+                "message": "What's the weather in San Francisco?",
+            },
+        }
+{% else %}
+        data = {
+            "input": {
+                "input": {
+                    "messages": [
+                        {"type": "human", "content": "Hello, AI!"},
+                        {"type": "ai", "content": "Hello!"},
+                        {"type": "human", "content": "How are you?"},
+                    ]
+                },
+                "config": {
+                    "metadata": {"user_id": "test-user", "session_id": "test-session"}
+                },
+            }
+        }
+{% endif %}
+        start_time = time.time()
+        with self.client.post(
+            url_path,
+            headers=headers,
+            json=data,
+            catch_response=True,
+{%- if cookiecutter.is_adk %}
+            name="/streamQuery async_stream_query",
+{%- else %}
+            name="/stream_messages first message",
+{%- endif %}
+            stream=True,
+            params={"alt": "sse"},
+        ) as response:
+            if response.status_code == 200:
+                events = []
+                for line in response.iter_lines():
+                    if line:
+                        line_str = line.decode("utf-8")
+                        events.append(line_str)
+
+                        if "429 Too Many Requests" in line_str:
+                            self.environment.events.request.fire(
+                                request_type="POST",
+                                name=f"{url_path} rate_limited 429s",
+                                response_time=0,
+                                response_length=len(line),
+                                response=response,
+                                context={},
+                            )
+                end_time = time.time()
+                total_time = end_time - start_time
+                self.environment.events.request.fire(
+                    request_type="POST",
+{%- if cookiecutter.is_adk %}
+                    name="/streamQuery end",
+{%- else %}
+                    name="/stream_messages end",
+{%- endif %}
+                    response_time=total_time * 1000,  # Convert to milliseconds
+                    response_length=len(events),
+                    response=response,
+                    context={},
+                )
+            else:
+                response.failure(f"Unexpected status code: {response.status_code}")
+{%- endif %}

{src → agent_starter_pack}/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 (
@@ -235,6 +239,16 @@ class AgentEngineApp:
     default=None,
     help="Service account email to use for the agent engine",
 )
+@click.option(
+    "--staging-bucket-uri",
+    default=None,
+    help="GCS bucket URI for staging files (defaults to gs://{project}-agent-engine)",
+)
+@click.option(
+    "--artifacts-bucket-name",
+    default=None,
+    help="GCS bucket name for artifacts (defaults to gs://{project}-agent-engine)",
+)
 def deploy_agent_engine_app(
     project: str | None,
     location: str,
@@ -243,13 +257,31 @@ def deploy_agent_engine_app(
     extra_packages: tuple[str, ...],
     set_env_vars: str | None,
     service_account: str | None,
+    staging_bucket_uri: str | None,
+    artifacts_bucket_name: str | None,
 ) -> AgentEngine:
     """Deploy the agent engine app to Vertex AI."""
+
+    logging.basicConfig(level=logging.INFO)
+
     # Parse environment variables if provided
     env_vars = parse_env_vars(set_env_vars)
 
     if not project:
         _, project = google.auth.default()
+    if not staging_bucket_uri:
+        staging_bucket_uri = f"gs://{project}-agent-engine"
+    if not artifacts_bucket_name:
+        artifacts_bucket_name = f"gs://{project}-agent-engine"
+
+{%- if "adk" in cookiecutter.tags %}
+    create_bucket_if_not_exists(
+        bucket_name=artifacts_bucket_name, project=project, location=location
+    )
+{%- endif %}
+    create_bucket_if_not_exists(
+        bucket_name=staging_bucket_uri, project=project, location=location
+    )
 
     print("""
     ╔═══════════════════════════════════════════════════════════╗
@@ -259,18 +291,7 @@ def deploy_agent_engine_app(
     ╚═══════════════════════════════════════════════════════════╝
     """)
 
-    logging.basicConfig(level=logging.INFO)
     extra_packages_list = list(extra_packages)
-    staging_bucket_uri = f"gs://{project}-agent-engine"
-{%- if "adk" in cookiecutter.tags %}
-    artifacts_bucket_name = f"{project}-{{cookiecutter.project_name}}-logs"
-    create_bucket_if_not_exists(
-        bucket_name=artifacts_bucket_name, project=project, location=location
-    )
-{%- endif %}
-    create_bucket_if_not_exists(
-        bucket_name=staging_bucket_uri, project=project, location=location
-    )
 
     # Initialize vertexai client
     client = vertexai.Client(
@@ -282,7 +303,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 +335,11 @@ def deploy_agent_engine_app(
         requirements=requirements,
         staging_bucket=staging_bucket_uri,
         labels=labels,
+        gcs_dir_name=agent_name,
+{%- if cookiecutter.is_adk_live %}
+        agent_server_mode=AgentServerMode.EXPERIMENTAL,  # Enable bidi streaming
+        resource_limits={"cpu": "4", "memory": "8Gi"},
+{%- endif %}
     )
 
     agent_config = {

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

@@ -75,11 +75,20 @@ 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 %}
+{%- if cookiecutter.is_adk_live %}
+    print("\n✅ Deployment successful! Run your agent with: `make playground-remote`")
+{%- else %}
     print(
-        f"\n✅ Deployment successful! Test your agent: notebooks/adk_app_testing.ipynb"
-        f"\n📊 View in console: {console_url}\n"
+        "\n✅ Deployment successful! Test your agent: notebooks/adk_app_testing.ipynb"
     )
+{%- endif %}
+    service_account = remote_agent.api_resource.spec.service_account
+    print(f"Service Account: {service_account}")
+    print(f"\n📊 View in console: {console_url}\n")
 {%- else %}
-    print(f"\n✅ Deployment successful!\n📊 View in console: {console_url}\n")
+    print("\n✅ Deployment successful!")
+    service_account = remote_agent.api_resource.spec.service_account
+    print(f"Service Account: {service_account}")
+    print(f"\n📊 View in console: {console_url}\n")
 {%- endif %}