agent-starter-pack 0.2.3__py3-none-any.whl → 0.3.1__py3-none-any.whl

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

@@ -54,7 +54,14 @@ 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 %}
+        data = {
+            "input": {
+                "message": "What's the weather in San Francisco?",
+                "user_id": "test",
+            }
+        }
+{% else %}
         data = {
             "input": {
                 "input": {
@@ -69,7 +76,7 @@ class ChatStreamUser(HttpUser):
                 },
             }
         }
-
+{% endif %}
         start_time = time.time()
         with self.client.post(
             url_path,

src/deployment_targets/cloud_run/app/server.py

@@ -11,7 +11,32 @@
 # 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 "adk" in cookiecutter.tags %}
+import os
+
+from fastapi import FastAPI
+from google.adk.cli.fast_api import get_fast_api_app
+from google.cloud import logging as google_cloud_logging
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider, export
+
+from app.utils.tracing import CloudTraceLoggingSpanExporter
+from app.utils.typing import Feedback
+
+logging_client = google_cloud_logging.Client()
+logger = logging_client.logger(__name__)
+
+provider = TracerProvider()
+processor = export.BatchSpanProcessor(CloudTraceLoggingSpanExporter())
+provider.add_span_processor(processor)
+trace.set_tracer_provider(provider)
 
+AGENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+app: FastAPI = get_fast_api_app(agent_dir=AGENT_DIR, web=False)
+
+app.title = "{{cookiecutter.project_name}}"
+app.description = "API for interacting with the Agent {{cookiecutter.project_name}}"
+{%- else %}
 import logging
 import os
 from collections.abc import Generator
@@ -40,7 +65,7 @@ try:
         app_name=app.title,
         disable_batch=False,
         exporter=CloudTraceLoggingSpanExporter(),
-        instruments={% raw %}{{% endraw %}{%- for instrumentation in cookiecutter.otel_instrumentations %}{{ instrumentation }}{% if not loop.last %}, {% endif %}{%- endfor %}{% raw %}}{% endraw %},
+        instruments={Instruments.LANGCHAIN, Instruments.CREW},
     )
 except Exception as e:
     logging.error("Failed to initialize Telemetry: %s", str(e))
@@ -91,20 +116,6 @@ def redirect_root_to_docs() -> RedirectResponse:
     return RedirectResponse(url="/docs")
 
 
-@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.post("/stream_messages")
 def stream_chat_events(request: Request) -> StreamingResponse:
     """Stream chat events in response to an input request.
@@ -119,6 +130,21 @@ def stream_chat_events(request: Request) -> StreamingResponse:
         stream_messages(input=request.input, config=request.config),
         media_type="text/event-stream",
     )
+{%- endif %}
+
+
+@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"}
 
 
 # Main execution

src/deployment_targets/cloud_run/tests/integration/test_server_e2e.py

@@ -32,7 +32,11 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 BASE_URL = "http://127.0.0.1:8000/"
+{%- if "adk" in cookiecutter.tags %}
+STREAM_URL = BASE_URL + "run_sse"
+{%- else %}
 STREAM_URL = BASE_URL + "stream_messages"
+{%- endif %}
 FEEDBACK_URL = BASE_URL + "feedback"
 
 HEADERS = {"Content-Type": "application/json"}
@@ -116,23 +120,72 @@ 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 %}
+    # Create session first
+    user_id = "user_123"
+    session_id = "session_abc"
+    session_data = {"state": {"preferred_language": "English", "visit_count": 5}}
+    session_response = requests.post(
+        f"{BASE_URL}/apps/app/users/{user_id}/sessions/{session_id}",
+        headers=HEADERS,
+        json=session_data,
+        timeout=10,
+    )
+    assert session_response.status_code == 200
 
+    # Then send chat message
+    data = {
+        "app_name": "app",
+        "user_id": user_id,
+        "session_id": session_id,
+        "new_message": {
+            "role": "user",
+            "parts": [{"text": "What's the weather in San Francisco?"}],
+        },
+        "streaming": True,
+    }
+{% else %}
     data = {
         "input": {
             "messages": [
                 {"type": "human", "content": "Hello, AI!"},
                 {"type": "ai", "content": "Hello!"},
-                {"type": "human", "content": "What is the weather in NY?"},
+                {"type": "human", "content": "Who are you?"},
             ]
         },
         "config": {"metadata": {"user_id": "test-user", "session_id": "test-session"}},
     }
-
+{% endif %}
     response = requests.post(
         STREAM_URL, headers=HEADERS, json=data, stream=True, timeout=10
     )
     assert response.status_code == 200
 
+{%- if "adk" in cookiecutter.tags %}
+    # Parse SSE events from response
+    events = []
+    for line in response.iter_lines():
+        if line:
+            # SSE format is "data: {json}"
+            line_str = line.decode("utf-8")
+            if line_str.startswith("data: "):
+                event_json = line_str[6:]  # Remove "data: " prefix
+                event = json.loads(event_json)
+                events.append(event)
+
+    assert events, "No events received from stream"
+    # Check for valid content in the response
+    has_text_content = False
+    for event in events:
+        content = event.get("content")
+        if (
+            content is not None
+            and content.get("parts")
+            and any(part.get("text") for part in content["parts"])
+        ):
+            has_text_content = True
+            break
+{%- else %}
     events = [json.loads(line) for line in response.iter_lines() if line]
     assert events, "No events received from stream"
 
@@ -155,12 +208,12 @@ def test_chat_stream(server_fixture: subprocess.Popen[str]) -> None:
             has_content = True
             break
     assert has_content, "At least one message should have content"
+{%- endif %}
 
 
 def test_chat_stream_error_handling(server_fixture: subprocess.Popen[str]) -> None:
     """Test the chat stream error handling."""
     logger.info("Starting chat stream error handling test")
-
     data = {
         "input": {"messages": [{"type": "invalid_type", "content": "Cause an error"}]}
     }
@@ -182,7 +235,11 @@ def test_collect_feedback(server_fixture: subprocess.Popen[str]) -> None:
     # Create sample feedback data
     feedback_data = {
         "score": 4,
+{%- if "adk" in cookiecutter.tags %}
+        "invocation_id": str(uuid.uuid4()),
+{%- else %}
         "run_id": str(uuid.uuid4()),
+{%- endif %}
         "text": "Great response!",
     }
 

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

@@ -28,7 +28,7 @@ Trigger the Locust load test with the following command:
 locust -f tests/load_test/load_test.py \
 -H http://127.0.0.1:8000 \
 --headless \
--t 30s -u 60 -r 2 \
+-t 30s -u 10 -r 2 \
 --csv=tests/load_test/.results/results \
 --html=tests/load_test/.results/report.html
 ```

src/deployment_targets/cloud_run/tests/load_test/load_test.py

@@ -15,9 +15,20 @@
 import json
 import os
 import time
+{%- if "adk" in cookiecutter.tags %}
+import uuid
 
+import requests
 from locust import HttpUser, between, task
+{%- else %}
 
+from locust import HttpUser, between, task
+{%- endif %}
+{% if "adk" in cookiecutter.tags %}
+ENDPOINT = "/run_sse"
+{% else %}
+ENDPOINT = "/stream_messages"
+{% endif %}
 
 class ChatStreamUser(HttpUser):
     """Simulates a user interacting with the chat stream API."""
@@ -30,7 +41,30 @@ class ChatStreamUser(HttpUser):
         headers = {"Content-Type": "application/json"}
         if os.environ.get("_ID_TOKEN"):
             headers["Authorization"] = f"Bearer {os.environ['_ID_TOKEN']}"
+{%- if "adk" in cookiecutter.tags %}
+        # Create session first
+        user_id = f"user_{uuid.uuid4()}"
+        session_id = f"session_{uuid.uuid4()}"
+        session_data = {"state": {"preferred_language": "English", "visit_count": 5}}
+        requests.post(
+            f"{self.client.base_url}/apps/app/users/{user_id}/sessions/{session_id}",
+            headers=headers,
+            json=session_data,
+            timeout=10,
+        )
 
+        # Send chat message
+        data = {
+            "app_name": "app",
+            "user_id": user_id,
+            "session_id": session_id,
+            "new_message": {
+                "role": "user",
+                "parts": [{"text": "What's the weather in San Francisco?"}],
+            },
+            "streaming": True,
+        }
+{%- else %}
         data = {
             "input": {
                 "messages": [
@@ -43,43 +77,42 @@ class ChatStreamUser(HttpUser):
                 "metadata": {"user_id": "test-user", "session_id": "test-session"}
             },
         }
-
+{%- endif %}
         start_time = time.time()
 
         with self.client.post(
-            "/stream_messages",
+            ENDPOINT,
+            name=f"{ENDPOINT} message",
             headers=headers,
             json=data,
             catch_response=True,
-            name="/stream_messages first message",
             stream=True,
+            params={"alt": "sse"},
         ) as response:
             if response.status_code == 200:
                 events = []
                 for line in response.iter_lines():
                     if line:
+{%- if "adk" in cookiecutter.tags %}
+                        # SSE format is "data: {json}"
+                        line_str = line.decode("utf-8")
+                        if line_str.startswith("data: "):
+                            event_json = line_str[6:]  # Remove "data: " prefix
+                            event = json.loads(event_json)
+                            events.append(event)
+{%- else %}
                         event = json.loads(line)
                         events.append(event)
-                        for chunk in event:
-                            if (
-                                isinstance(chunk, dict)
-                                and chunk.get("type") == "constructor"
-                            ):
-                                if not chunk.get("kwargs", {}).get("content"):
-                                    continue
-                                response.success()
-                                end_time = time.time()
-                                total_time = end_time - start_time
-                                self.environment.events.request.fire(
-                                    request_type="POST",
-                                    name="/stream_messages end",
-                                    response_time=total_time
-                                    * 1000,  # Convert to milliseconds
-                                    response_length=len(json.dumps(events)),
-                                    response=response,
-                                    context={},
-                                )
-                                return
-                response.failure("No valid response content received")
+{%- endif %}
+                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(json.dumps(events)),
+                    response=response,
+                    context={},
+                )
             else:
                 response.failure(f"Unexpected status code: {response.status_code}")

src/frontends/live_api_react/frontend/package-lock.json

@@ -2027,9 +2027,9 @@
       }
     },
     "node_modules/@babel/runtime": {
-      "version": "7.26.0",
-      "resolved": "https://registry.npmjs.org/@babel/runtime/-/runtime-7.26.0.tgz",
-      "integrity": "sha512-FDSOghenHTiToteC/QRlv2q3DhPZ/oOXTBoirfWNx1Cx3TMVcGWQtMMmQcSvb/JjpNeGzx8Pq/b4fKEJuWm1sw==",
+      "version": "7.27.0",
+      "resolved": "https://registry.npmjs.org/@babel/runtime/-/runtime-7.27.0.tgz",
+      "integrity": "sha512-VtPOkrdPHZsKc/clNqyi9WUA8TINkZ4cGk63UUE3u4pmB2k+ZMQRDuIOagv8UVd6j7k0T3+RRIb7beKTebNbcw==",
       "license": "MIT",
       "dependencies": {
         "regenerator-runtime": "^0.14.0"

src/frontends/streamlit_adk/frontend/side_bar.py

@@ -0,0 +1,214 @@
+# 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.
+
+# ruff: noqa: RUF015
+import json
+import os
+import uuid
+from typing import Any
+
+from frontend.utils.chat_utils import save_chat
+from frontend.utils.multimodal_utils import (
+    HELP_GCS_CHECKBOX,
+    HELP_MESSAGE_MULTIMODALITY,
+    upload_files_to_gcs,
+)
+
+EMPTY_CHAT_NAME = "Empty chat"
+NUM_CHAT_IN_RECENT = 3
+DEFAULT_BASE_URL = "http://localhost:8000/"
+
+DEFAULT_REMOTE_AGENT_ENGINE_ID = "N/A"
+if os.path.exists("deployment_metadata.json"):
+    with open("deployment_metadata.json") as f:
+        DEFAULT_REMOTE_AGENT_ENGINE_ID = json.load(f)["remote_agent_engine_id"]
+DEFAULT_AGENT_CALLABLE_PATH = "app.agent_engine_app.AgentEngineApp"
+
+
+class SideBar:
+    """Manages the sidebar components of the Streamlit application."""
+
+    def __init__(self, st: Any) -> None:
+        """
+        Initialize the SideBar.
+
+        Args:
+            st (Any): The Streamlit object for rendering UI components.
+        """
+        self.st = st
+
+    def init_side_bar(self) -> None:
+        """Initialize and render the sidebar components."""
+        with self.st.sidebar:
+            default_agent_type = (
+                "Remote URL" if os.path.exists("Dockerfile") else "Local Agent"
+            )
+            use_agent_path = self.st.selectbox(
+                "Select Agent Type",
+                ["Local Agent", "Remote Agent Engine ID", "Remote URL"],
+                index=["Local Agent", "Remote Agent Engine ID", "Remote URL"].index(
+                    default_agent_type
+                ),
+                help="'Local Agent' uses a local implementation, 'Remote Agent Engine ID' connects to a deployed Vertex AI agent, and 'Remote URL' connects to a custom endpoint.",
+            )
+
+            if use_agent_path == "Local Agent":
+                self.agent_callable_path = self.st.text_input(
+                    label="Agent Callable Path",
+                    value=os.environ.get(
+                        "AGENT_CALLABLE_PATH", DEFAULT_AGENT_CALLABLE_PATH
+                    ),
+                )
+                self.remote_agent_engine_id = None
+                self.url_input_field = None
+                self.should_authenticate_request = False
+            elif use_agent_path == "Remote Agent Engine ID":
+                self.remote_agent_engine_id = self.st.text_input(
+                    label="Remote Agent Engine ID",
+                    value=os.environ.get(
+                        "REMOTE_AGENT_ENGINE_ID", DEFAULT_REMOTE_AGENT_ENGINE_ID
+                    ),
+                )
+                self.agent_callable_path = None
+                self.url_input_field = None
+                self.should_authenticate_request = False
+            else:
+                self.url_input_field = self.st.text_input(
+                    label="Service URL",
+                    value=os.environ.get("SERVICE_URL", DEFAULT_BASE_URL),
+                )
+                self.should_authenticate_request = self.st.checkbox(
+                    label="Authenticate request",
+                    value=False,
+                    help="If checked, any request to the server will contain an"
+                    "Identity token to allow authentication. "
+                    "See the Cloud Run documentation to know more about authentication:"
+                    "https://cloud.google.com/run/docs/authenticating/service-to-service",
+                )
+                self.agent_callable_path = None
+                self.remote_agent_engine_id = None
+
+            col1, col2, col3 = self.st.columns(3)
+            with col1:
+                if self.st.button("+ New chat"):
+                    if (
+                        len(
+                            self.st.session_state.user_chats[
+                                self.st.session_state["session_id"]
+                            ]["messages"]
+                        )
+                        > 0
+                    ):
+                        self.st.session_state.invocation_id = None
+
+                        self.st.session_state["session_id"] = str(uuid.uuid4())
+                        self.st.session_state.session_db.get_session(
+                            session_id=self.st.session_state["session_id"],
+                        )
+                        self.st.session_state.user_chats[
+                            self.st.session_state["session_id"]
+                        ] = {
+                            "title": EMPTY_CHAT_NAME,
+                            "messages": [],
+                        }
+
+            with col2:
+                if self.st.button("Delete chat"):
+                    self.st.session_state.invocation_id = None
+                    self.st.session_state.session_db.clear()
+                    self.st.session_state.user_chats.pop(
+                        self.st.session_state["session_id"]
+                    )
+                    if len(self.st.session_state.user_chats) > 0:
+                        chat_id = list(self.st.session_state.user_chats.keys())[0]
+                        self.st.session_state["session_id"] = chat_id
+                        self.st.session_state.session_db.get_session(
+                            session_id=self.st.session_state["session_id"],
+                        )
+                    else:
+                        self.st.session_state["session_id"] = str(uuid.uuid4())
+                        self.st.session_state.user_chats[
+                            self.st.session_state["session_id"]
+                        ] = {
+                            "title": EMPTY_CHAT_NAME,
+                            "messages": [],
+                        }
+            with col3:
+                if self.st.button("Save chat"):
+                    save_chat(self.st)
+
+            self.st.subheader("Recent")  # Style the heading
+
+            all_chats = list(reversed(self.st.session_state.user_chats.items()))
+            for chat_id, chat in all_chats[:NUM_CHAT_IN_RECENT]:
+                if self.st.button(chat["title"], key=chat_id):
+                    self.st.session_state.invocation_id = None
+                    self.st.session_state["session_id"] = chat_id
+                    self.st.session_state.session_db.get_session(
+                        session_id=self.st.session_state["session_id"],
+                    )
+
+            with self.st.expander("Other chats"):
+                for chat_id, chat in all_chats[NUM_CHAT_IN_RECENT:]:
+                    if self.st.button(chat["title"], key=chat_id):
+                        self.st.session_state.invocation_id = None
+                        self.st.session_state["session_id"] = chat_id
+                        self.st.session_state.session_db.get_session(
+                            session_id=self.st.session_state["session_id"],
+                        )
+
+            self.st.divider()
+            self.st.header("Upload files from local")
+            bucket_name = self.st.text_input(
+                label="GCS Bucket for upload",
+                value=os.environ.get("BUCKET_NAME", "gs://your-bucket-name"),
+            )
+            if "checkbox_state" not in self.st.session_state:
+                self.st.session_state.checkbox_state = True
+
+            self.st.session_state.checkbox_state = self.st.checkbox(
+                "Upload to GCS first (suggested)", value=False, help=HELP_GCS_CHECKBOX
+            )
+
+            self.uploaded_files = self.st.file_uploader(
+                label="Send files from local",
+                accept_multiple_files=True,
+                key=f"uploader_images_{self.st.session_state.uploader_key}",
+                type=[
+                    "png",
+                    "jpg",
+                    "jpeg",
+                    "txt",
+                    "docx",
+                    "pdf",
+                    "rtf",
+                    "csv",
+                    "tsv",
+                    "xlsx",
+                ],
+            )
+            if self.uploaded_files and self.st.session_state.checkbox_state:
+                upload_files_to_gcs(self.st, bucket_name, self.uploaded_files)
+
+            self.st.divider()
+
+            self.st.header("Upload files from GCS")
+            self.gcs_uris = self.st.text_area(
+                "GCS uris (comma-separated)",
+                value=self.st.session_state["gcs_uris_to_be_sent"],
+                key=f"upload_text_area_{self.st.session_state.uploader_key}",
+                help=HELP_MESSAGE_MULTIMODALITY,
+            )
+
+            self.st.caption(f"Note: {HELP_MESSAGE_MULTIMODALITY}")