agent-starter-pack 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

src/deployment_targets/agent_engine/app/agent_engine_app.py

@@ -12,6 +12,53 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# mypy: disable-error-code="attr-defined"
+{%- if "adk" in cookiecutter.tags %}
+import datetime
+import json
+import logging
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import google.auth
+import vertexai
+from google.cloud import logging as google_cloud_logging
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider, export
+from vertexai import agent_engines
+from vertexai.preview.reasoning_engines import AdkApp
+
+from app.agent import root_agent
+from app.utils.gcs import create_bucket_if_not_exists
+from app.utils.tracing import CloudTraceLoggingSpanExporter
+from app.utils.typing import Feedback
+
+
+class AgentEngineApp(AdkApp):
+    def set_up(self) -> None:
+        """Set up logging and tracing for the agent engine app."""
+        super().set_up()
+        logging_client = google_cloud_logging.Client()
+        self.logger = logging_client.logger(__name__)
+        provider = TracerProvider()
+        processor = export.BatchSpanProcessor(CloudTraceLoggingSpanExporter())
+        provider.add_span_processor(processor)
+        trace.set_tracer_provider(provider)
+
+    def register_feedback(self, feedback: dict[str, Any]) -> None:
+        """Collect and log feedback."""
+        feedback_obj = Feedback.model_validate(feedback)
+        self.logger.log_struct(feedback_obj.model_dump(), severity="INFO")
+
+    def register_operations(self) -> Mapping[str, Sequence]:
+        """Registers the operations of the Agent.
+
+        Extends the base operations to include feedback registration functionality.
+        """
+        operations = super().register_operations()
+        operations[""] = operations[""] + ["register_feedback"]
+        return operations
+{%- else %}
 import datetime
 import json
 import logging
@@ -26,16 +73,12 @@ import vertexai
 from google.cloud import logging as google_cloud_logging
 from langchain_core.runnables import RunnableConfig
 from traceloop.sdk import Instruments, Traceloop
-from vertexai.preview import reasoning_engines
+from vertexai import agent_engines
 
 from app.utils.gcs import create_bucket_if_not_exists
 from app.utils.tracing import CloudTraceLoggingSpanExporter
 from app.utils.typing import Feedback, InputChat, dumpd, ensure_valid_config
 
-logging.basicConfig(
-    level=logging.INFO,
-)
-
 
 class AgentEngineApp:
     """Class for managing agent engine functionality."""
@@ -70,7 +113,6 @@ class AgentEngineApp:
             )
         except Exception as e:
             logging.error("Failed to initialize Telemetry: %s", str(e))
-
         self.runnable = agent
 
     # Add any additional variables here that should be included in the tracing logs
@@ -111,11 +153,6 @@ class AgentEngineApp:
             dumped_chunk = dumpd(chunk)
             yield dumped_chunk
 
-    def register_feedback(self, feedback: dict[str, Any]) -> None:
-        """Collect and log feedback."""
-        feedback_obj = Feedback.model_validate(feedback)
-        self.logger.log_struct(feedback_obj.model_dump(), severity="INFO")
-
     def query(
         self,
         *,
@@ -124,8 +161,15 @@ class AgentEngineApp:
         **kwargs: Any,
     ) -> Any:
         """Process a single input and return the agent's response."""
+        config = ensure_valid_config(config)
+        self.set_tracing_properties(config=config)
         return dumpd(self.runnable.invoke(input=input, config=config, **kwargs))
 
+    def register_feedback(self, feedback: dict[str, Any]) -> None:
+        """Collect and log feedback."""
+        feedback_obj = Feedback.model_validate(feedback)
+        self.logger.log_struct(feedback_obj.model_dump(), severity="INFO")
+
     def register_operations(self) -> Mapping[str, Sequence]:
         """Registers the operations of the Agent.
 
@@ -142,6 +186,7 @@ class AgentEngineApp:
             "": ["query", "register_feedback"],
             "stream": ["stream_query"],
         }
+{%- endif %}
 
 
 def deploy_agent_engine_app(
@@ -151,7 +196,7 @@ def deploy_agent_engine_app(
     requirements_file: str = ".requirements.txt",
     extra_packages: list[str] = ["./app"],
     env_vars: dict[str, str] | None = None,
-) -> reasoning_engines.ReasoningEngine:
+) -> agent_engines.AgentEngine:
     """Deploy the agent engine app to Vertex AI."""
 
     staging_bucket = f"gs://{project}-agent-engine"
@@ -164,24 +209,25 @@ def deploy_agent_engine_app(
     # Read requirements
     with open(requirements_file) as f:
         requirements = f.read().strip().split("\n")
-
-    agent = AgentEngineApp(project_id=project, env_vars=env_vars)
-
+{% if "adk" in cookiecutter.tags %}
+    agent_engine = AgentEngineApp(
+        agent=root_agent, env_vars=env_vars, enable_tracing=True
+    )
+{% else %}
+    agent_engine = AgentEngineApp(project_id=project, env_vars=env_vars)
+{% endif %}
     # Common configuration for both create and update operations
     agent_config = {
-        "reasoning_engine": agent,
+        "agent_engine": agent_engine,
         "display_name": agent_name,
-        "description": "This is a sample custom application in Agent Engine that uses LangGraph",
+        "description": "{{cookiecutter.agent_description}}",
         "extra_packages": extra_packages,
     }
     logging.info(f"Agent config: {agent_config}")
     agent_config["requirements"] = requirements
 
     # Check if an agent with this name already exists
-    existing_agents = reasoning_engines.ReasoningEngine.list(
-        filter=f"display_name={agent_name}"
-    )
-
+    existing_agents = list(agent_engines.list(filter=f"display_name={agent_name}"))
     if existing_agents:
         # Update the existing agent with new configuration
         logging.info(f"Updating existing agent: {agent_name}")
@@ -189,7 +235,7 @@ def deploy_agent_engine_app(
     else:
         # Create a new agent if none exists
         logging.info(f"Creating new agent: {agent_name}")
-        remote_agent = reasoning_engines.ReasoningEngine.create(**agent_config)
+        remote_agent = agent_engines.create(**agent_config)
 
     config = {
         "remote_agent_engine_id": remote_agent.resource_name,

src/deployment_targets/agent_engine/app/utils/gcs.py

@@ -14,8 +14,8 @@
 
 import logging
 
+import google.cloud.storage as storage
 from google.api_core import exceptions
-from google.cloud import storage
 
 
 def create_bucket_if_not_exists(bucket_name: str, project: str, location: str) -> None:

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

@@ -15,18 +15,80 @@
 import logging
 
 import pytest
+{%- if "adk" in cookiecutter.tags %}
+from google.adk.events.event import Event
 
+from app.agent import root_agent
 from app.agent_engine_app import AgentEngineApp
+{%- else %}
+
+from app.agent_engine_app import AgentEngineApp
+{%- endif %}
 
 
 @pytest.fixture
 def agent_app() -> AgentEngineApp:
     """Fixture to create and set up AgentEngineApp instance"""
+{%- if "adk" in cookiecutter.tags %}
+    app = AgentEngineApp(agent=root_agent)
+{%- else %}
     app = AgentEngineApp()
+{%- endif %}
     app.set_up()
     return app
 
+{% if "adk" in cookiecutter.tags %}
+def test_agent_stream_query(agent_app: AgentEngineApp) -> None:
+    """
+    Integration test for the agent stream query functionality.
+    Tests that the agent returns valid streaming responses.
+    """
+    # Create message and events for the stream_query
+    message = "What's the weather in San Francisco?"
+    events = list(agent_app.stream_query(message=message, user_id="test"))
+    assert len(events) > 0, "Expected at least one chunk in response"
+
+    # Check for valid content in the response
+    has_text_content = False
+    for event in events:
+        validated_event = Event.model_validate(event)
+        content = validated_event.content
+        if (
+            content is not None
+            and content.parts
+            and any(part.text for part in content.parts)
+        ):
+            has_text_content = True
+            break
+
+    assert has_text_content, "Expected at least one event with text content"
 
+
+def test_agent_feedback(agent_app: AgentEngineApp) -> None:
+    """
+    Integration test for the agent feedback functionality.
+    Tests that feedback can be registered successfully.
+    """
+    feedback_data = {
+        "score": 5,
+        "text": "Great response!",
+        "invocation_id": "test-run-123",
+    }
+
+    # Should not raise any exceptions
+    agent_app.register_feedback(feedback_data)
+
+    # Test invalid feedback
+    with pytest.raises(ValueError):
+        invalid_feedback = {
+            "score": "invalid",  # Score must be numeric
+            "text": "Bad feedback",
+            "invocation_id": "test-run-123",
+        }
+        agent_app.register_feedback(invalid_feedback)
+
+    logging.info("All assertions passed for agent feedback test")
+{% else %}
 def test_agent_stream_query(agent_app: AgentEngineApp) -> None:
     """
     Integration test for the agent stream query functionality.
@@ -118,3 +180,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 %}

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/Dockerfile

@@ -14,7 +14,7 @@
 
 FROM python:3.11-slim
 
-RUN pip install --no-cache-dir uv
+RUN pip install --no-cache-dir uv==0.6.12
 
 WORKDIR /code
 

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/frontend/utils/stream_handler.py

@@ -28,7 +28,7 @@ import streamlit as st
 import vertexai
 from google.auth.exceptions import DefaultCredentialsError
 from langchain_core.messages import AIMessage, ToolMessage
-from vertexai.preview import reasoning_engines
+from vertexai import agent_engines
 
 from frontend.utils.multimodal_utils import format_content
 
@@ -43,7 +43,7 @@ def get_remote_agent(remote_agent_engine_id: str) -> Any:
     project_id = parts[1]
     location = parts[3]
     vertexai.init(project=project_id, location=location)
-    return reasoning_engines.ReasoningEngine(remote_agent_engine_id)
+    return agent_engines.AgentEngine(remote_agent_engine_id)
 
 
 @st.cache_resource
@@ -159,7 +159,7 @@ class Client:
             if self.authenticate_request:
                 headers["Authorization"] = f"Bearer {self.id_token}"
             with requests.post(
-                self.url, json=data, headers=headers, stream=True, timeout=10
+                self.url, json=data, headers=headers, stream=True, timeout=60
             ) as response:
                 for line in response.iter_lines():
                     if line: