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

{src → agent_starter_pack}/deployment_targets/cloud_run/tests/integration/test_server_e2e.py

@@ -11,6 +11,207 @@
 # 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 asyncio
+import json
+import logging
+import subprocess
+import sys
+import threading
+import time
+from collections.abc import Iterator
+from typing import Any
+
+import pytest
+import requests
+from websockets.asyncio.client import connect
+
+# Configure logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+WS_URL = "ws://127.0.0.1:8000/ws"
+FEEDBACK_URL = "http://127.0.0.1:8000/feedback"
+
+
+def log_output(pipe: Any, log_func: Any) -> None:
+    """Log the output from the given pipe."""
+    for line in iter(pipe.readline, ""):
+        log_func(line.strip())
+
+
+def start_server() -> subprocess.Popen[str]:
+    """Start the server in local mode."""
+    command = [
+        sys.executable,
+        "-m",
+        "uvicorn",
+        "{{cookiecutter.agent_directory}}.server:app",
+        "--host",
+        "0.0.0.0",
+        "--port",
+        "8000",
+    ]
+    process = subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        bufsize=1,
+        encoding="utf-8",
+    )
+
+    # Start threads to log stdout and stderr in real-time
+    threading.Thread(
+        target=log_output, args=(process.stdout, logger.info), daemon=True
+    ).start()
+    threading.Thread(
+        target=log_output, args=(process.stderr, logger.error), daemon=True
+    ).start()
+
+    return process
+
+
+def wait_for_server(timeout: int = 60, interval: int = 1) -> bool:
+    """Wait for the server to be ready."""
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            response = requests.get("http://127.0.0.1:8000/docs", timeout=10)
+            if response.status_code == 200:
+                logger.info("Server is ready")
+                return True
+        except Exception:
+            pass
+        time.sleep(interval)
+    logger.error(f"Server did not become ready within {timeout} seconds")
+    return False
+
+
+@pytest.fixture(scope="module")
+def server_fixture(request: Any) -> Iterator[subprocess.Popen[str]]:
+    """Pytest fixture to start and stop the server for testing."""
+    logger.info("Starting server process")
+    server_process = start_server()
+    if not wait_for_server():
+        pytest.fail("Server failed to start")
+    logger.info("Server process started")
+
+    def stop_server() -> None:
+        logger.info("Stopping server process")
+        server_process.terminate()
+        try:
+            server_process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            logger.warning("Server process did not terminate, killing it")
+            server_process.kill()
+            server_process.wait()
+        logger.info("Server process stopped")
+
+    request.addfinalizer(stop_server)
+    yield server_process
+
+
+@pytest.mark.asyncio
+async def test_websocket_audio_input(server_fixture: subprocess.Popen[str]) -> None:
+    """Test websocket with audio input in local mode."""
+
+    async def send_message(websocket: Any, message: dict[str, Any]) -> None:
+        """Helper to send JSON messages."""
+        await websocket.send(json.dumps(message))
+
+    async def receive_message(websocket: Any, timeout: float = 5.0) -> dict[str, Any]:
+        """Helper to receive messages with timeout."""
+        try:
+            response = await asyncio.wait_for(websocket.recv(), timeout=timeout)
+            if isinstance(response, bytes):
+                return json.loads(response.decode())
+            if isinstance(response, str):
+                return json.loads(response)
+            return response
+        except asyncio.TimeoutError as exc:
+            raise TimeoutError(
+                f"No response received within {timeout} seconds"
+            ) from exc
+
+    try:
+        await asyncio.sleep(2)
+
+        async with connect(WS_URL, ping_timeout=10, close_timeout=10) as websocket:
+            try:
+                # Wait for setupComplete
+                setup_response = await receive_message(websocket, timeout=10.0)
+                assert "setupComplete" in setup_response
+                logger.info("Received setupComplete")
+
+                # Send dummy audio chunk with user_id
+                dummy_audio = bytes([0] * 1024)
+                audio_msg = {
+                    "user_id": "test-user",
+                    "realtimeInput": {
+                        "mediaChunks": [
+                            {
+                                "mimeType": "audio/pcm;rate=16000",
+                                "data": dummy_audio.hex(),
+                            }
+                        ]
+                    },
+                }
+                await send_message(websocket, audio_msg)
+                logger.info("Sent audio chunk")
+
+                # Send text message to complete the turn (matching frontend format)
+                text_msg = {
+                    "content": {
+                        "role": "user",
+                        "parts": [{"text": "Test audio"}],
+                    }
+                }
+                await send_message(websocket, text_msg)
+                logger.info("Sent text completion")
+
+                # Collect responses
+                responses = []
+                for _ in range(10):
+                    try:
+                        response = await receive_message(websocket, timeout=5.0)
+                        responses.append(response)
+                        logger.info(f"Received: {response}")
+
+                        if isinstance(response, dict) and response.get("turn_complete"):
+                            break
+                    except TimeoutError:
+                        break
+
+                # Verify we got responses
+                assert len(responses) > 0, "No responses received"
+
+                logger.info(f"Audio test passed. Received {len(responses)} responses")
+
+            finally:
+                await websocket.close()
+
+    except Exception as e:
+        logger.error(f"Audio test failed: {e}")
+        raise
+
+
+def test_feedback_endpoint(server_fixture: subprocess.Popen[str]) -> None:
+    """Test the feedback endpoint."""
+    feedback_data = {
+        "score": 5,
+        "text": "Great response!",
+        "invocation_id": "test-run-123",
+        "user_id": "test-user",
+        "log_type": "feedback",
+    }
+
+    response = requests.post(FEEDBACK_URL, json=feedback_data, timeout=10)
+    assert response.status_code == 200
+    assert response.json() == {"status": "success"}
+    logger.info("Feedback endpoint test passed")
+{%- else %}
 
 import json
 import logging
@@ -32,7 +233,7 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 BASE_URL = "http://127.0.0.1:8000/"
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
 STREAM_URL = BASE_URL + "run_sse"
 {%- else %}
 STREAM_URL = BASE_URL + "stream_messages"
@@ -86,7 +287,7 @@ def start_server() -> subprocess.Popen[str]:
     return process
 
 
-def wait_for_server(timeout: int = 60, interval: int = 1) -> bool:
+def wait_for_server(timeout: int = 90, interval: int = 1) -> bool:
     """Wait for the server to be ready."""
     start_time = time.time()
     while time.time() - start_time < timeout:
@@ -124,7 +325,7 @@ def server_fixture(request: Any) -> Iterator[subprocess.Popen[str]]:
 def test_chat_stream(server_fixture: subprocess.Popen[str]) -> None:
     """Test the chat stream functionality."""
     logger.info("Starting chat stream test")
-{% if "adk" in cookiecutter.tags %}
+{% if cookiecutter.is_adk %}
     # Create session first
     user_id = "test_user_123"
     session_data = {"state": {"preferred_language": "English", "visit_count": 1}}
@@ -168,7 +369,7 @@ def test_chat_stream(server_fixture: subprocess.Popen[str]) -> None:
     )
     assert response.status_code == 200
 
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
     # Parse SSE events from response
     events = []
     for line in response.iter_lines():
@@ -242,7 +443,7 @@ def test_collect_feedback(server_fixture: subprocess.Popen[str]) -> None:
     # Create sample feedback data
     feedback_data = {
         "score": 4,
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
         "invocation_id": str(uuid.uuid4()),
 {%- else %}
         "run_id": str(uuid.uuid4()),
@@ -285,3 +486,4 @@ def cleanup_agent_engine_sessions() -> None:
     except Exception as e:
         logger.warning(f"Failed to cleanup agent engine sessions: {e}")
 {%- endif %}
+{%- endif %}

{src → agent_starter_pack}/deployment_targets/cloud_run/tests/load_test/README.md

@@ -1,6 +1,87 @@
 # Robust Load Testing for Generative AI Applications
 
 This directory provides a comprehensive load testing framework for your Generative AI application, leveraging the power of [Locust](http://locust.io), a leading open-source load testing tool.
+{%- if cookiecutter.agent_name == "adk_live" %}
+
+## Local Load Testing
+
+Follow these steps to execute load tests on your local machine:
+
+**1. Start the FastAPI Server:**
+
+Launch the FastAPI server in a separate terminal:
+
+```bash
+uv run uvicorn {{cookiecutter.agent_directory}}.server:app --host 0.0.0.0 --port 8000 --reload
+```
+
+**2. (In another tab) Create virtual environment with Locust**
+Using another terminal tab, This is suggested to avoid conflicts with the existing application python environment.
+
+```bash
+python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1 websockets
+```
+
+**3. Execute the Load Test:**
+Trigger the Locust load test with the following command:
+
+```bash
+locust -f tests/load_test/load_test.py \
+-H http://127.0.0.1:8000 \
+--headless \
+-t 30s -u 2 -r 2 \
+--csv=tests/load_test/.results/results \
+--html=tests/load_test/.results/report.html
+```
+
+This command initiates a 30-second load test, simulating 2 users spawning per second, reaching a maximum of 60 concurrent users.
+
+**Results:**
+
+Comprehensive CSV and HTML reports detailing the load test performance will be generated and saved in the `tests/load_test/.results` directory.
+
+## Remote Load Testing (Targeting Cloud Run)
+
+This framework also supports load testing against remote targets, such as a staging Cloud Run instance. This process is seamlessly integrated into the Continuous Delivery (CD) pipeline.
+
+**Prerequisites:**
+
+- **Dependencies:** Ensure your environment has the same dependencies required for local testing.
+- **Cloud Run Invoker Role:** You'll need the `roles/run.invoker` role to invoke the Cloud Run service.
+
+**Steps:**
+
+**1. Start Cloud Run Proxy:**
+
+Start the proxy in a separate terminal to expose your Cloud Run service on localhost. The proxy automatically handles IAM authentication:
+
+```bash
+gcloud run services proxy YOUR_SERVICE_NAME --port=8080 --region us-central1 --quiet
+```
+
+Replace `YOUR_SERVICE_NAME` with your Cloud Run service name. The `--quiet` flag auto-approves component installation prompts. You can optionally specify `--tag` to target a specific traffic tag.
+
+**2. (In another tab) Create virtual environment with Locust:**
+
+Using another terminal tab:
+
+```bash
+python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1 websockets
+```
+
+**3. Execute the Load Test:**
+
+Execute load tests against the proxied service. The proxy handles authentication automatically:
+
+```bash
+locust -f tests/load_test/load_test.py \
+-H http://127.0.0.1:8080 \
+--headless \
+-t 30s -u 2 -r 2 \
+--csv=tests/load_test/.results/results \
+--html=tests/load_test/.results/report.html
+```
+{%- else %}
 
 ## Local Load Testing
 
@@ -81,3 +162,4 @@ locust -f tests/load_test/load_test.py \
 --csv=tests/load_test/.results/results \
 --html=tests/load_test/.results/report.html
 ```
+{%- endif %}

agent_starter_pack/deployment_targets/cloud_run/tests/load_test/load_test.py

@@ -0,0 +1,249 @@
+# 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 typing import Any
+
+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: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.ws_url = (
+            self.host.replace("http://", "ws://").replace("https://", "wss://") + "/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 os
+import time
+{%- if cookiecutter.is_adk %}
+import uuid
+
+import requests
+from locust import HttpUser, between, task
+{%- else %}
+
+from locust import HttpUser, between, task
+{%- endif %}
+{% if cookiecutter.is_adk %}
+ENDPOINT = "/run_sse"
+{% else %}
+ENDPOINT = "/stream_messages"
+{% endif %}
+
+class ChatStreamUser(HttpUser):
+    """Simulates a user interacting with the chat stream API."""
+
+    wait_time = between(1, 3)  # Wait 1-3 seconds between tasks
+
+    @task
+    def chat_stream(self) -> None:
+        """Simulates a chat stream interaction."""
+        headers = {"Content-Type": "application/json"}
+        if os.environ.get("_ID_TOKEN"):
+            headers["Authorization"] = f"Bearer {os.environ['_ID_TOKEN']}"
+{%- if cookiecutter.is_adk %}
+        # Create session first
+        user_id = f"user_{uuid.uuid4()}"
+        session_data = {"state": {"preferred_language": "English", "visit_count": 1}}
+
+        session_url = f"{self.client.base_url}/apps/{{cookiecutter.agent_directory}}/users/{user_id}/sessions"
+        session_response = requests.post(
+            session_url,
+            headers=headers,
+            json=session_data,
+            timeout=10,
+        )
+
+        # Get session_id from response
+        session_id = session_response.json()["id"]
+
+        # Send chat message
+        data = {
+            "app_name": "{{cookiecutter.agent_directory}}",
+            "user_id": user_id,
+            "session_id": session_id,
+            "new_message": {
+                "role": "user",
+                "parts": [{"text": "Hello! Weather in New york?"}],
+            },
+            "streaming": True,
+        }
+{%- else %}
+        data = {
+            "input": {
+                "messages": [
+                    {"type": "human", "content": "Hello, AI!"},
+                    {"type": "ai", "content": "Hello!"},
+                    {"type": "human", "content": "Who are you?"},
+                ]
+            },
+            "config": {
+                "metadata": {"user_id": "test-user", "session_id": "test-session"}
+            },
+        }
+{%- endif %}
+        start_time = time.time()
+
+        with self.client.post(
+            ENDPOINT,
+            name=f"{ENDPOINT} message",
+            headers=headers,
+            json=data,
+            catch_response=True,
+            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"{ENDPOINT} 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",
+                    name=f"{ENDPOINT} end",
+                    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 %}