agent-starter-pack 0.0.1b0__py3-none-any.whl

agents/multimodal_live_api/tests/integration/test_server_e2e.py

@@ -0,0 +1,254 @@
+# 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
+import os
+import subprocess
+import sys
+import threading
+import time
+import uuid
+from collections.abc import Iterator
+from typing import Any
+
+import pytest
+import requests
+import websockets.client
+
+# Configure logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+BASE_URL = "ws://127.0.0.1:8000/"
+WS_URL = BASE_URL + "ws"
+
+FEEDBACK_URL = "http://127.0.0.1:8000/feedback"
+HEADERS = {"Content-Type": "application/json"}
+
+
+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 FastAPI server using subprocess and log its output."""
+    command = [
+        sys.executable,
+        "-m",
+        "uvicorn",
+        "app.server:app",
+        "--host",
+        "0.0.0.0",
+        "--port",
+        "8000",
+    ]
+    env = os.environ.copy()
+    env["INTEGRATION_TEST"] = "TRUE"
+    process = subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        bufsize=1,
+        env=env,
+    )
+
+    # 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="session")
+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()
+        server_process.wait()
+        logger.info("Server process stopped")
+
+    request.addfinalizer(stop_server)
+    yield server_process
+
+
+@pytest.mark.asyncio
+async def test_websocket_connection(server_fixture: subprocess.Popen[str]) -> None:
+    """Test the websocket connection and message exchange."""
+
+    async def send_message(websocket: Any, message: dict[str, Any]) -> None:
+        """Helper function to send messages and log them."""
+        await websocket.send(json.dumps(message))
+
+    async def receive_message(websocket: Any, timeout: float = 5.0) -> dict[str, Any]:
+        """Helper function 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 websockets.connect(
+            WS_URL, ping_timeout=10, close_timeout=10
+        ) as websocket:
+            try:
+                # Wait for initial ready message
+                initial_response = None
+                for _ in range(10):
+                    try:
+                        initial_response = await receive_message(websocket, timeout=5.0)
+                        if (
+                            initial_response is not None
+                            and initial_response.get("status")
+                            == "Backend is ready for conversation"
+                        ):
+                            break
+                    except TimeoutError:
+                        if _ == 9:
+                            raise
+                        continue
+
+                assert (
+                    initial_response is not None
+                    and initial_response.get("status")
+                    == "Backend is ready for conversation"
+                )
+
+                # Send messages
+                setup_msg = {"setup": {"run_id": "test-run", "user_id": "test-user"}}
+                await send_message(websocket, setup_msg)
+
+                dummy_audio = bytes([0] * 1024)
+                audio_msg = {
+                    "realtimeInput": {
+                        "mediaChunks": [
+                            {
+                                "mimeType": "audio/pcm;rate=16000",
+                                "data": dummy_audio.hex(),
+                            }
+                        ]
+                    }
+                }
+                await send_message(websocket, audio_msg)
+
+                text_msg = {
+                    "clientContent": {
+                        "turns": [
+                            {"role": "user", "parts": [{"text": "Hello, how are you?"}]}
+                        ],
+                        "turnComplete": True,
+                    }
+                }
+                await send_message(websocket, text_msg)
+
+                # Collect responses with timeout
+                responses = []
+                try:
+                    while True:
+                        try:
+                            response = await receive_message(websocket, timeout=10.0)
+                            responses.append(response)
+                            if (
+                                len(responses) >= 3
+                            ):  # Exit after receiving enough responses
+                                break
+                        except TimeoutError:
+                            break
+                except asyncio.TimeoutError:
+                    logger.info("Response collection timed out")
+
+                # Verify responses
+                assert len(responses) > 0, "No responses received from server"
+                assert any(
+                    isinstance(r, dict) and "serverContent" in r for r in responses
+                )
+                logger.info(
+                    f"Test completed successfully. Received {len(responses)} responses"
+                )
+
+            finally:
+                # Ensure websocket is closed properly
+                await websocket.close()
+
+    except Exception as e:
+        logger.error(f"Test failed with error: {e!s}")
+        raise
+
+    finally:
+        # Clean up any remaining tasks
+        tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
+        for task in tasks:
+            task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+
+def test_collect_feedback(server_fixture: subprocess.Popen[str]) -> None:
+    """
+    Test the feedback collection endpoint (/feedback) to ensure it properly
+    logs the received feedback.
+    """
+    # Create sample feedback data
+    feedback_data = {
+        "score": 4,
+        "text": "Great response!",
+        "run_id": str(uuid.uuid4()),
+        "user_id": "user1",
+        "log_type": "feedback",
+    }
+
+    response = requests.post(
+        FEEDBACK_URL, json=feedback_data, headers=HEADERS, timeout=10
+    )
+    assert response.status_code == 200

agents/multimodal_live_api/tests/load_test/load_test.py

@@ -0,0 +1,40 @@
+# 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 time
+
+from locust import HttpUser, between, task
+
+
+class DummyUser(HttpUser):
+    """Simulates a user for testing purposes."""
+
+    wait_time = between(1, 3)  # Wait 1-3 seconds between tasks
+
+    @task
+    def dummy_task(self) -> None:
+        """A dummy task that simulates work without making actual requests."""
+        # Simulate some processing time
+        time.sleep(0.1)
+
+        # Record a successful dummy request
+        self.environment.events.request.fire(
+            request_type="POST",
+            name="dummy_endpoint",
+            response_time=100,
+            response_length=1024,
+            response=None,
+            context={},
+            exception=None,
+        )

agents/multimodal_live_api/tests/unit/test_server.py

@@ -0,0 +1,143 @@
+# 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 json
+import logging
+import os
+from collections.abc import Generator
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from fastapi.testclient import TestClient
+from google.auth.credentials import Credentials
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@pytest.fixture(autouse=True)
+def mock_google_cloud_credentials() -> Generator[None, None, None]:
+    """Mock Google Cloud credentials for testing."""
+    with patch.dict(
+        os.environ,
+        {
+            "GOOGLE_APPLICATION_CREDENTIALS": "/path/to/mock/credentials.json",
+            "GOOGLE_CLOUD_PROJECT_ID": "mock-project-id",
+        },
+    ):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def mock_google_auth_default() -> Generator[None, None, None]:
+    """Mock the google.auth.default function for testing."""
+    mock_credentials = MagicMock(spec=Credentials)
+    mock_project = "mock-project-id"
+
+    with patch("google.auth.default", return_value=(mock_credentials, mock_project)):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def mock_dependencies() -> Generator[None, None, None]:
+    """
+    Mock Vertex AI dependencies for testing.
+    Patches genai client and tool functions.
+    """
+    with (
+        patch("app.server.genai_client") as mock_genai,
+        patch("app.server.tool_functions") as mock_tools,
+    ):
+        mock_genai.aio.live.connect = AsyncMock()
+        mock_tools.return_value = {}
+        yield
+
+
+@pytest.mark.asyncio
+async def test_websocket_endpoint() -> None:
+    """
+    Test the websocket endpoint to ensure it correctly handles
+    websocket connections and messages.
+    """
+    from app.server import app
+
+    mock_session = AsyncMock()
+    mock_session._ws = AsyncMock()
+    # Configure mock to return proper response format and close after one message
+    mock_session._ws.recv.side_effect = [
+        json.dumps(
+            {
+                "serverContent": {
+                    "modelTurn": {
+                        "role": "model",
+                        "parts": [{"text": "Hello, how can I help you?"}],
+                    }
+                }
+            }
+        ).encode(),  # Encode as bytes since recv(decode=False) is used
+        None,  # Add None to trigger StopAsyncIteration after first message
+    ]
+
+    with patch("app.server.genai_client") as mock_genai:
+        mock_genai.aio.live.connect.return_value.__aenter__.return_value = mock_session
+        client = TestClient(app)
+        with client.websocket_connect("/ws") as websocket:
+            # Test initial connection message
+            data = websocket.receive_json()
+            assert data["status"] == "Backend is ready for conversation"
+
+            # Test sending a message
+            websocket.send_json(
+                {"setup": {"run_id": "test-run", "user_id": "test-user"}}
+            )
+
+            # Test sending audio stream
+            dummy_audio = bytes([0] * 1024)  # 1KB of silence
+            websocket.send_json(
+                {
+                    "realtimeInput": {
+                        "mediaChunks": [
+                            {
+                                "mimeType": "audio/pcm;rate=16000",
+                                "data": dummy_audio.hex(),
+                            }
+                        ]
+                    }
+                }
+            )
+
+            # Receive response as bytes
+            response = websocket.receive_bytes()
+            response_data = json.loads(response.decode())
+            assert "serverContent" in response_data
+
+            # Verify mock interactions
+            mock_genai.aio.live.connect.assert_called_once()
+            assert mock_session._ws.recv.called
+
+
+@pytest.mark.asyncio
+async def test_websocket_error_handling() -> None:
+    """Test websocket error handling."""
+    from app.server import app
+
+    with patch("app.server.genai_client") as mock_genai:
+        mock_genai.aio.live.connect.side_effect = Exception("Connection failed")
+
+        client = TestClient(app)
+        with pytest.raises(Exception) as exc:
+            with client.websocket_connect("/ws"):
+                pass
+        assert str(exc.value) == "Connection failed"

src/base_template/.gitignore

@@ -0,0 +1,197 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*.pyc
+*$py.class
+**/dist
+/tmp
+/out-tsc
+/bazel-out
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+Pipfile.lock
+Pipfile
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+.venv*
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# macOS
+.DS_Store
+
+# PyCharm
+.idea
+
+# User-specific files
+.terraform*
+.Terraform*
+
+
+tmp*
+
+# Node
+**/node_modules
+npm-debug.log
+yarn-error.log
+
+# IDEs and editors
+.idea/
+.project
+.classpath
+.c9/
+*.launch
+.settings/
+*.sublime-workspace
+
+# Visual Studio Code
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+.history/*
+
+# Miscellaneous
+**/.angular/*
+/.angular/cache
+.sass-cache/
+/connect.lock
+/coverage
+/libpeerconnection.log
+testem.log
+/typings
+
+# System files
+.DS_Store
+Thumbs.db
+*.vscode*
+
+# Starter pack specific
+.persist_vector_store
+tests/load_test/.results/*.html
+tests/load_test/.results/*.csv
+locust_env
+my_env.tfvars
+.streamlit_chats
+.saved_chats
+.env
+.requirements.txt

src/base_template/Makefile

@@ -0,0 +1,37 @@
+test:
+	uv run pytest tests/unit && uv run pytest tests/integration
+
+install:
+	uv sync --dev {% if cookiecutter.agent_name != 'multimodal_live_api' %}--extra streamlit{%- endif %} --extra jupyter --frozen{% if cookiecutter.agent_name == 'multimodal_live_api' %} && npm --prefix frontend install{%- endif %}
+
+playground:
+{%- if cookiecutter.deployment_target == 'cloud_run' %}
+	uv run uvicorn app.server:app --host 0.0.0.0 --port 8000 --reload &
+{%- endif %}
+{%- if cookiecutter.agent_name == 'multimodal_live_api' %}
+	npm --prefix frontend start
+{%- else %}
+	{% if cookiecutter.deployment_target == 'agent_engine' %}PYTHONPATH=. {% endif %}uv run streamlit run frontend/streamlit_app.py --browser.serverAddress=localhost --server.enableCORS=false --server.enableXsrfProtection=false
+{%- endif %}
+
+backend:
+{%- if cookiecutter.deployment_target == 'cloud_run' %}
+	uv run uvicorn app.server:app --host 0.0.0.0 --port 8000 --reload
+{%- elif cookiecutter.deployment_target == 'agent_engine' %}
+	uv export --no-hashes --no-sources --no-header --no-emit-project --frozen > .requirements.txt && uv run app/agent_engine_app.py
+{%- endif %}
+
+{% if cookiecutter.deployment_target == 'cloud_run' -%}
+ui:
+{%- if cookiecutter.agent_name == 'multimodal_live_api' %}
+	npm --prefix frontend start
+{%- else %}
+	uv run streamlit run streamlit/streamlit_app.py --browser.serverAddress=localhost --server.enableCORS=false --server.enableXsrfProtection=false
+{%- endif %}
+{%- endif %}
+
+lint:
+	uv run codespell
+	uv run ruff check . --diff
+	uv run ruff format . --check --diff
+	uv run mypy .