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

src/base_template/{% if cookiecutter.cicd_runner == 'google_cloud_build' %}.cloudbuild{% else %}unused_.cloudbuild{% endif %}/deploy-to-prod.yaml

@@ -20,7 +20,7 @@ steps:
     args:
       - -c
       - |
-        cd data_ingestion && pip install uv==0.6.12 --user && cd data_ingestion_pipeline && \
+        cd data_ingestion && pip install uv==0.8.13 --user && cd data_ingestion_pipeline && \
         uv sync --locked && uv run python submit_pipeline.py
     env:
       - "PIPELINE_ROOT=${_PIPELINE_GCS_ROOT_PROD}"
@@ -63,7 +63,7 @@ steps:
     args:
       - "-c"
       - |
-        pip install uv==0.6.12 --user && uv sync --locked
+        pip install uv==0.8.13 --user && uv sync --locked
     env:
       - 'PATH=/usr/local/bin:/usr/bin:~/.local/bin'
 

src/base_template/{% if cookiecutter.cicd_runner == 'google_cloud_build' %}.cloudbuild{% else %}unused_.cloudbuild{% endif %}/pr_checks.yaml

@@ -20,7 +20,7 @@ steps:
     args:
       - "-c"
       - |
-        pip install uv==0.6.12 --user && uv sync --locked
+        pip install uv==0.8.13 --user && uv sync --locked
     env:
       - 'PATH=/usr/local/bin:/usr/bin:~/.local/bin'
 

src/base_template/{% if cookiecutter.cicd_runner == 'google_cloud_build' %}.cloudbuild{% else %}unused_.cloudbuild{% endif %}/staging.yaml

@@ -20,7 +20,7 @@ steps:
     args:
       - -c
       - |
-        cd data_ingestion && pip install uv==0.6.12 --user && cd data_ingestion_pipeline && \
+        cd data_ingestion && pip install uv==0.8.13 --user && cd data_ingestion_pipeline && \
         uv sync --locked && uv run python submit_pipeline.py
     env:
       - "PIPELINE_ROOT=${_PIPELINE_GCS_ROOT_STAGING}"
@@ -97,7 +97,7 @@ steps:
     args:
       - "-c"
       - |
-        pip install uv==0.6.12 --user && uv sync --locked
+        pip install uv==0.8.13 --user && uv sync --locked
     env:
       - 'PATH=/usr/local/bin:/usr/bin:~/.local/bin'
 
@@ -127,31 +127,113 @@ steps:
 {%- endif %}
 
   # Load Testing
+{%- if cookiecutter.deployment_target == 'cloud_run' and cookiecutter.agent_name == 'adk_live' %}
+  - name: "europe-west4-docker.pkg.dev/production-ai-template/starter-pack/e2e-tests"
+    id: load_test
+    entrypoint: /bin/bash
+    args:
+      - "-c"
+      - |
+        # Install load test dependencies
+        pip install locust==2.31.1 websockets
+
+        # Install Cloud Run proxy component
+        apt-get update && apt-get install -y google-cloud-cli-cloud-run-proxy
+
+        # Start Cloud Run proxy in background
+        gcloud run services proxy {{cookiecutter.project_name}} \
+        --port=8080 \
+        --region ${_REGION} \
+        --project ${_STAGING_PROJECT_ID} \
+        --quiet &
+        _PROXY_PID=$$!
+
+        # Wait for proxy to be ready
+        echo "Waiting for proxy to start..."
+        sleep 10
+
+        # Run load test and capture exit code
+        locust -f tests/load_test/load_test.py \
+        --headless \
+        -H http://127.0.0.1:8080 \
+        -t 30s -u 2 -r 2 \
+        --csv=tests/load_test/.results/results \
+        --html=tests/load_test/.results/report.html
+        _LOAD_TEST_EXIT_CODE=$$?
+
+        # Clean up proxy
+        kill $$_PROXY_PID || true
+
+        # Exit with load test result to fail build if tests failed
+        exit $$_LOAD_TEST_EXIT_CODE
+{%- elif cookiecutter.deployment_target == 'cloud_run' %}
   - name: "python:3.11-slim"
     id: load_test
     entrypoint: /bin/bash
     args:
       - "-c"
       - |
-{%- if cookiecutter.deployment_target == 'cloud_run' %}
         export _ID_TOKEN=$(cat id_token.txt)
         export _STAGING_URL=$(cat staging_url.txt)
-{%- elif cookiecutter.deployment_target == 'agent_engine' %}
-        export _AUTH_TOKEN=$(cat auth_token.txt)
-{%- endif %}
         pip install locust==2.31.1 --user
         locust -f tests/load_test/load_test.py \
         --headless \
-{%- if cookiecutter.deployment_target == 'cloud_run' %}
         -H $$_STAGING_URL \
         -t 30s -u 10 -r 0.5 \
+        --csv=tests/load_test/.results/results \
+        --html=tests/load_test/.results/report.html
+    env:
+      - 'PATH=/usr/local/bin:/usr/bin:~/.local/bin'
+{%- elif cookiecutter.deployment_target == 'agent_engine' and cookiecutter.is_adk_live %}
+  - name: "python:3.11-slim"
+    id: load_test
+    entrypoint: /bin/bash
+    args:
+      - "-c"
+      - |
+        # Start expose app in remote mode (uses deployment_metadata.json by default)
+        uv run python -m {{cookiecutter.agent_directory}}.utils.expose_app --mode remote &
+        EXPOSE_PID=$$!
+
+        # Wait for expose app to be ready
+        sleep 10
+
+        # Run load test against local expose app
+        uv run --with locust==2.31.1 --with websockets locust -f tests/load_test/load_test.py \
+        -H http://127.0.0.1:8000 \
+        --headless \
+        -t 30s -u 2 -r 1 \
+        --csv=tests/load_test/.results/results \
+        --html=tests/load_test/.results/report.html
+        LOCUST_EXIT_CODE=$$?
+
+        # Stop expose app
+        kill $$EXPOSE_PID
+
+        # Exit with error if locust test failed
+        if [ $$LOCUST_EXIT_CODE -ne 0 ]; then
+          echo "Load test failed with exit code $$LOCUST_EXIT_CODE"
+          exit $$LOCUST_EXIT_CODE
+        fi
+    env:
+      - 'PATH=/usr/local/bin:/usr/bin:~/.local/bin'
 {%- elif cookiecutter.deployment_target == 'agent_engine' %}
+  - name: "python:3.11-slim"
+    id: load_test
+    entrypoint: /bin/bash
+    args:
+      - "-c"
+      - |
+        export _AUTH_TOKEN=$(cat auth_token.txt)
+        pip install locust==2.31.1 --user
+        locust -f tests/load_test/load_test.py \
+        --headless \
         -t 30s -u 2 -r 0.5 \
-{%- endif %}
         --csv=tests/load_test/.results/results \
         --html=tests/load_test/.results/report.html
     env:
       - 'PATH=/usr/local/bin:/usr/bin:~/.local/bin'
+{%- endif %}
 
   # Export Load Test Results to GCS
   - name: gcr.io/cloud-builders/gcloud

src/base_template/{{cookiecutter.agent_directory}}/utils/tracing.py

@@ -86,7 +86,7 @@ class CloudTraceLoggingSpanExporter(CloudTraceSpanExporter):
                 print(span_dict)
 
             # Log the span data to Google Cloud Logging
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
             self.logger.log_struct(
                 span_dict,
                 labels={

src/base_template/{{cookiecutter.agent_directory}}/utils/typing.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
 {%- if cookiecutter.deployment_target == 'cloud_run' %}
 import uuid
 from typing import (
@@ -57,7 +57,7 @@ from pydantic import (
 {%- endif %}
 
 
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
 {%- if cookiecutter.deployment_target == 'cloud_run' %}
 
 
@@ -103,7 +103,7 @@ class Feedback(BaseModel):
 
     score: int | float
     text: str | None = ""
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
     invocation_id: str
 {%- else %}
     run_id: str
@@ -111,7 +111,7 @@ class Feedback(BaseModel):
     log_type: Literal["feedback"] = "feedback"
     service_name: Literal["{{cookiecutter.project_name}}"] = "{{cookiecutter.project_name}}"
     user_id: str = ""
-{% if "adk" not in cookiecutter.tags %}
+{% if not cookiecutter.is_adk %}
 
 def ensure_valid_config(config: RunnableConfig | None) -> RunnableConfig:
     """Ensures a valid RunnableConfig by setting defaults for missing fields."""

src/cli/commands/create.py

@@ -935,7 +935,7 @@ def set_gcp_project(project_id: str, set_quota_project: bool = True) -> None:
                 capture_output=True,
                 text=True,
             )
-        except subprocess.CalledProcessError:
+        except subprocess.CalledProcessError as e:
             logging.debug(f"Setting quota project failed: {e.stderr}")
 
     console.print(f"> Successfully configured project: {project_id}")

src/cli/utils/template.py

@@ -112,7 +112,7 @@ def get_available_agents(deployment_target: str | None = None) -> dict:
     # Define priority agents that should appear first
     PRIORITY_AGENTS = [
         "adk_base",
-        "adk_gemini_fullstack",
+        "adk_live",
         "agentic_rag",
         "langgraph_base_react",
     ]
@@ -740,6 +740,8 @@ def process_template(
                 ),
                 "settings": settings,
                 "tags": tags,
+                "is_adk": "adk" in tags,
+                "is_adk_live": "adk_live" in tags,
                 "deployment_target": deployment_target or "",
                 "cicd_runner": cicd_runner or "google_cloud_build",
                 "session_type": session_type or "",
@@ -756,7 +758,12 @@ def process_template(
                 "_copy_without_render": [
                     "*.ipynb",  # Don't render notebooks
                     "*.json",  # Don't render JSON files
-                    "frontend/*",  # Don't render frontend directory
+                    "*.tsx",  # Don't render TypeScript React files
+                    "*.ts",  # Don't render TypeScript files
+                    "*.jsx",  # Don't render JavaScript React files
+                    "*.js",  # Don't render JavaScript files
+                    "*.css",  # Don't render CSS files
+                    "frontend/**/*",  # Don't render frontend directory recursively
                     "notebooks/*",  # Don't render notebooks directory
                     ".git/*",  # Don't render git directory
                     "__pycache__/*",  # Don't render cache
@@ -1109,12 +1116,12 @@ def should_exclude_path(
     path: pathlib.Path, agent_name: str, agent_directory: str = "app"
 ) -> bool:
     """Determine if a path should be excluded based on the agent type."""
-    if agent_name == "live_api":
-        # Exclude the unit test utils folder and agent utils folder for live_api
+    if agent_name == "adk_live":
+        # Exclude the unit test utils folder and agent utils folder for adk_live
         if "tests/unit/test_utils" in str(path) or f"{agent_directory}/utils" in str(
             path
         ):
-            logging.debug(f"Excluding path for live_api: {path}")
+            logging.debug(f"Excluding path for adk_live: {path}")
             return True
     return False
 

src/deployment_targets/agent_engine/tests/integration/test_agent_engine_app.py

@@ -11,11 +11,212 @@
 # 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
-{%- if "adk" in cookiecutter.tags %}
+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 using expose_app in local mode."""
+    command = [
+        sys.executable,
+        "-m",
+        "uvicorn",
+        "app.utils.expose_app: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!",
+        "run_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 logging
+
+import pytest
+{%- if cookiecutter.is_adk %}
 from google.adk.events.event import Event
 
 from {{cookiecutter.agent_directory}}.agent import root_agent
@@ -29,7 +230,7 @@ from {{cookiecutter.agent_directory}}.agent_engine_app import AgentEngineApp
 @pytest.fixture
 def agent_app() -> AgentEngineApp:
     """Fixture to create and set up AgentEngineApp instance"""
-{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.is_adk %}
     app = AgentEngineApp(agent=root_agent)
 {%- else %}
     app = AgentEngineApp()
@@ -37,7 +238,7 @@ def agent_app() -> AgentEngineApp:
     app.set_up()
     return app
 
-{% if "adk" in cookiecutter.tags %}
+{% if cookiecutter.is_adk %}
 @pytest.mark.asyncio
 async def test_agent_stream_query(agent_app: AgentEngineApp) -> None:
     """
@@ -183,4 +384,4 @@ def test_agent_feedback(agent_app: AgentEngineApp) -> None:
         agent_app.register_feedback(invalid_feedback)
 
     logging.info("All assertions passed for agent feedback test")
-{% endif %}
+{% endif %}{% endif %}

src/deployment_targets/agent_engine/tests/load_test/README.md

@@ -1,3 +1,49 @@
+{%- if cookiecutter.agent_name == "adk_live" %}
+# WebSocket Load Testing for Remote Agent Engine
+
+This directory provides a comprehensive load testing framework for your Agent Engine application using WebSocket connections, leveraging the power of [Locust](http://locust.io), a leading open-source load testing tool.
+
+The load test simulates realistic user interactions by:
+- Establishing WebSocket connections
+- Sending audio chunks in the proper `realtimeInput` format
+- Sending text messages to complete turns
+- Collecting and measuring responses until `turn_complete`
+
+## Load Testing with Remote Agent Engine
+
+**1. Start the Expose App in Remote Mode:**
+
+Launch the expose app server in a separate terminal, pointing to your deployed agent engine:
+
+```bash
+uv run python -m app.utils.expose_app --mode remote --remote-id <your-agent-engine-id>
+```
+
+Or if you have `deployment_metadata.json` in your project root:
+
+```bash
+uv run python -m app.utils.expose_app --mode remote
+```
+
+**2. Execute the Load Test:**
+
+Using another terminal tab, trigger the Locust load test:
+
+```bash
+uv run --with locust==2.31.1 --with websockets locust -f tests/load_test/load_test.py \
+-H http://127.0.0.1:8000 \
+--headless \
+-t 30s -u 1 -r 1 \
+--csv=tests/load_test/.results/results \
+--html=tests/load_test/.results/report.html
+```
+
+This command initiates a 30-second load test with 1 concurrent user.
+
+**Results:**
+
+Comprehensive CSV and HTML reports detailing the load test performance will be generated and saved in the `tests/load_test/.results` directory.
+{%- else %}
 # 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.
@@ -34,4 +80,5 @@ Follow these steps to execute load tests:
    ```
 
    This command initiates a 30-second load test, simulating 2 users spawning per second, reaching a maximum of 10 concurrent users.
+{%- endif %}