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

agents/langgraph_base_react/template/.templateconfig.yaml

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-description: "A agent implementing a base ReAct agent using LangGraph"
+description: "An agent implementing a base ReAct agent using LangGraph"
 settings:
   requires_data_ingestion: false
   deployment_targets: ["agent_engine", "cloud_run"]

src/base_template/Makefile

@@ -6,6 +6,14 @@ test:
 	uv run pytest tests/unit && uv run pytest tests/integration
 
 playground:
+{%- if "adk" in cookiecutter.tags %}
+	@echo "+-----------------------------------------------------------------------------+"
+	@echo "| Starting ADK Web Server via 'adk web' command.                              |"
+	@echo "|                                                                             |"
+	@echo "| Please select the app folder to start interacting with your agent.          |"
+	@echo "+-----------------------------------------------------------------------------+"
+	uv run adk web
+{%- else %}
 {%- if cookiecutter.deployment_target == 'cloud_run' %}
 	uv run uvicorn app.server:app --host 0.0.0.0 --port 8000 --reload &
 {%- endif %}
@@ -14,6 +22,7 @@ playground:
 {%- 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 %}
+{%- endif %}
 
 backend:
 {%- if cookiecutter.deployment_target == 'cloud_run' %}

src/base_template/README.md

@@ -47,7 +47,7 @@ make install && make playground
 | -------------------- | ------------------------------------------------------------------------------------------- |
 | `make install`       | Install all required dependencies using uv                                                  |
 {%- if cookiecutter.deployment_target == 'cloud_run' %}
-| `make playground`    | Launch local development environment with backend and frontend |
+| `make playground`    | Launch local development environment with backend and frontend{%- if "adk" in cookiecutter.tags %} - leveraging `adk web` command. {%- endif %}|
 | `make backend`       | Start backend server only |
 | `make ui`            | Launch Streamlit frontend without local backend |
 {%- elif cookiecutter.deployment_target == 'agent_engine' %}

src/base_template/app/__init__.py

@@ -0,0 +1,3 @@
+from app.agent import root_agent
+
+__all__ = ["root_agent"]

src/base_template/app/utils/tracing.py

@@ -17,8 +17,8 @@ import logging
 from collections.abc import Sequence
 from typing import Any
 
+import google.cloud.storage as storage
 from google.cloud import logging as google_cloud_logging
-from google.cloud import storage
 from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExportResult
@@ -86,8 +86,18 @@ class CloudTraceLoggingSpanExporter(CloudTraceSpanExporter):
                 print(span_dict)
 
             # Log the span data to Google Cloud Logging
+{%- if "adk" in cookiecutter.tags %}
+            self.logger.log_struct(
+                span_dict,
+                labels={
+                    "type": "agent_telemetry",
+                    "service_name": "{{cookiecutter.project_name}}",
+                },
+                severity="INFO",
+            )
+{%- else %}
             self.logger.log_struct(span_dict, severity="INFO")
-
+{%- endif %}
         # Export spans to Google Cloud Trace using the parent class method
         return super().export(spans)
 

src/base_template/app/utils/typing.py

@@ -12,6 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.deployment_target == 'cloud_run' %}
+import uuid
+from typing import (
+    Literal,
+)
+
+from google.adk.events.event import Event
+from google.genai.types import Content
+from pydantic import (
+    BaseModel,
+    Field,
+)
+{%- else %}
+from typing import (
+    Literal,
+)
+
+from pydantic import (
+    BaseModel,
+)
+{%- endif %}
+{%- else %}
 import json
 import uuid
 from typing import (
@@ -31,6 +54,24 @@ from pydantic import (
     BaseModel,
     Field,
 )
+{%- endif %}
+
+
+{%- if "adk" in cookiecutter.tags %}
+{%- if cookiecutter.deployment_target == 'cloud_run' %}
+
+
+class Request(BaseModel):
+    """Represents the input for a chat request with optional configuration."""
+
+    message: Content
+    events: list[Event]
+    user_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    session_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+
+    model_config = {"extra": "allow"}
+{%- endif %}
+{%- else %}
 
 
 class InputChat(BaseModel):
@@ -42,7 +83,7 @@ class InputChat(BaseModel):
         ..., description="The chat messages representing the current conversation."
     )
 
-{% if cookiecutter.deployment_target == 'cloud_run' %}
+
 class Request(BaseModel):
     """Represents the input for a chat request with optional configuration.
 
@@ -54,16 +95,23 @@ class Request(BaseModel):
     input: InputChat
     config: RunnableConfig | None = None
 
-{% endif %}
+{%- endif %}
+
+
 class Feedback(BaseModel):
     """Represents feedback for a conversation."""
 
     score: int | float
     text: str | None = ""
+{%- if "adk" in cookiecutter.tags %}
+    invocation_id: str
+{%- else %}
     run_id: str
+{%- endif %}
     log_type: Literal["feedback"] = "feedback"
     service_name: Literal["{{cookiecutter.project_name}}"] = "{{cookiecutter.project_name}}"
-
+    user_id: str = ""
+{% if "adk" not in cookiecutter.tags %}
 
 def ensure_valid_config(config: RunnableConfig | None) -> RunnableConfig:
     """Ensures a valid RunnableConfig by setting defaults for missing fields."""
@@ -99,7 +147,8 @@ def dumps(obj: Any) -> str:
         JSON string representation of the object
     """
     return json.dumps(obj, default=default_serialization)
-{% if cookiecutter.deployment_target == 'agent_engine' %}
+{%- if cookiecutter.deployment_target == 'agent_engine' %}
+
 
 def dumpd(obj: Any) -> Any:
     """
@@ -113,4 +162,5 @@ def dumpd(obj: Any) -> Any:
         Dict/list representation of the object that can be JSON serialized
     """
     return json.loads(dumps(obj))
+{%- endif %}
 {% endif %}

src/base_template/deployment/terraform/dev/variables.tf

@@ -32,7 +32,11 @@ variable "region" {
 variable "telemetry_logs_filter" {
   type        = string
   description = "Log Sink filter for capturing telemetry data. Captures logs with the `traceloop.association.properties.log_type` attribute set to `tracing`."
+{%- if "adk" in cookiecutter.tags %}
+  default     = "labels.service_name=\"{{cookiecutter.project_name}}\" labels.type=\"agent_telemetry\""
+{%- else %}
   default     = "jsonPayload.attributes.\"traceloop.association.properties.log_type\"=\"tracing\" jsonPayload.resource.attributes.\"service.name\"=\"{{cookiecutter.project_name}}\""
+{%- endif %}
 }
 
 variable "feedback_logs_filter" {

src/base_template/deployment/terraform/dev/vars/env.tfvars

@@ -7,9 +7,6 @@ dev_project_id = "your-dev-project-id"
 # The Google Cloud region you will use to deploy the infrastructure
 region = "us-central1"
 
-telemetry_logs_filter = "jsonPayload.attributes.\"traceloop.association.properties.log_type\"=\"tracing\" jsonPayload.resource.attributes.\"service.name\"=\"{{cookiecutter.project_name}}\""
-feedback_logs_filter = "jsonPayload.log_type=\"feedback\""
-
 {%- if cookiecutter.data_ingestion %}
 {%- if cookiecutter.datastore_type == "vertex_ai_search" %}
 # The value can only be one of "global", "us" and "eu".

src/base_template/deployment/terraform/variables.tf

@@ -52,7 +52,11 @@ variable "repository_name" {
 variable "telemetry_logs_filter" {
   type        = string
   description = "Log Sink filter for capturing telemetry data. Captures logs with the `traceloop.association.properties.log_type` attribute set to `tracing`."
+{%- if "adk" in cookiecutter.tags %}
+  default     = "labels.service_name=\"{{cookiecutter.project_name}}\" labels.type=\"agent_telemetry\""
+{%- else %}
   default     = "jsonPayload.attributes.\"traceloop.association.properties.log_type\"=\"tracing\" jsonPayload.resource.attributes.\"service.name\"=\"{{cookiecutter.project_name}}\""
+{%- endif %}
 }
 
 variable "feedback_logs_filter" {

src/base_template/deployment/terraform/vars/env.tfvars

@@ -19,10 +19,6 @@ repository_name = "repo-{{cookiecutter.project_name}}"
 # The Google Cloud region you will use to deploy the infrastructure
 region = "us-central1"
 
-telemetry_logs_filter = "jsonPayload.attributes.\"traceloop.association.properties.log_type\"=\"tracing\" jsonPayload.resource.attributes.\"service.name\"=\"{{cookiecutter.project_name}}\""
-
-feedback_logs_filter = "jsonPayload.log_type=\"feedback\""
-
 {%- if cookiecutter.data_ingestion %}
 pipeline_cron_schedule = "0 0 * * 0"
 

src/base_template/pyproject.toml

@@ -9,16 +9,18 @@ dependencies = [
 {%- for dep in cookiecutter.extra_dependencies %}
     "{{ dep }}",
 {%- endfor %}
-    "langchain-core~=0.3.9",
     "opentelemetry-exporter-gcp-trace~=1.9.0",
+{%- if "adk" not in cookiecutter.tags %}
+    "langchain-core~=0.3.9",
     "traceloop-sdk~=0.38.7",
+{%- endif %}
     "google-cloud-logging~=3.11.4",
 {%- if cookiecutter.deployment_target == 'cloud_run' %}
-    "google-cloud-aiplatform[evaluation]~=1.87.0",
+    "google-cloud-aiplatform[evaluation]~=1.88.0",
     "fastapi~=0.115.8",
     "uvicorn~=0.34.0"
 {%- elif cookiecutter.deployment_target == 'agent_engine' %}
-    "google-cloud-aiplatform[evaluation,reasoningengine]~=1.87.0"
+    "google-cloud-aiplatform[evaluation,agent-engines]~=1.88.0"
 {%- endif %}
 ]
 {% if cookiecutter.deployment_target == 'cloud_run' %}

src/{deployment_targets/agent_engine → base_template}/tests/unit/test_dummy.py

@@ -11,9 +11,10 @@
 # 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.
-
 """
 You can add your unit tests here.
+This is where you test your business logic, including agent functionality,
+data processing, and other core components of your application.
 """
 
 

src/cli/commands/create.py

@@ -306,10 +306,10 @@ def create(
                 f"This agent uses a datastore for grounded responses.\n"
                 f"The agent will work without data, but for optimal results:\n"
                 f"1. Set up dev environment:\n"
-                f"   [white italic]`export PROJECT_ID={project_id} && cd {cd_path} && make setup-dev-env`[/white italic]\n\n"
+                f"   [white italic]export PROJECT_ID={project_id} && cd {cd_path} && make setup-dev-env[/white italic]\n\n"
                 f"   See deployment/README.md for more info\n"
                 f"2. Run the data ingestion pipeline:\n"
-                f"   [white italic]`export PROJECT_ID={project_id} && cd {cd_path} && make data-ingestion`[/white italic]\n\n"
+                f"   [white italic]export PROJECT_ID={project_id} && cd {cd_path} && make data-ingestion[/white italic]\n\n"
                 f"   See data_ingestion/README.md for more info\n"
                 f"[bold white]=================================[/bold white]\n"
             )
@@ -647,6 +647,14 @@ def replace_region_in_files(
                     "_DATA_STORE_REGION: us", f"_DATA_STORE_REGION: {data_store_region}"
                 )
                 modified = True
+            elif '"DATA_STORE_REGION", "us"' in content:
+                if debug:
+                    logging.debug(f"Replacing DATA_STORE_REGION in {file_path}")
+                content = content.replace(
+                    '"DATA_STORE_REGION", "us"',
+                    f'"DATA_STORE_REGION", "{data_store_region}"',
+                )
+                modified = True
 
             if modified:
                 file_path.write_text(content)

src/cli/commands/setup_cicd.py

@@ -815,6 +815,9 @@ def setup_cicd(
 
         repo_url = f"https://github.com/{github_username}/{config.repository_name}"
         cloud_build_url = f"https://console.cloud.google.com/cloud-build/builds?project={config.cicd_project_id}"
+        # Sleep to allow resources to propagate
+        console.print("\n⏳ Waiting for resources to propagate...")
+        time.sleep(10)
 
         # Print final summary
         print_cicd_summary(config, github_username, repo_url, cloud_build_url)

src/cli/utils/gcp.py

@@ -45,7 +45,7 @@ def get_dummy_request(project_id: str, location: str) -> CountTokensRequest:
     """Creates a simple test request for Gemini."""
     return CountTokensRequest(
         contents=[{"role": "user", "parts": [{"text": "Hi"}]}],
-        endpoint=f"projects/{project_id}/locations/{location}/publishers/google/models/gemini-1.5-flash-002",
+        endpoint=f"projects/{project_id}/locations/{location}/publishers/google/models/gemini-2.0-flash",
     )
 
 

src/cli/utils/template.py

@@ -29,6 +29,9 @@ from src.cli.utils.version import get_current_version
 
 from .datastores import DATASTORES
 
+ADK_FILES = ["app/__init__.py"]
+NON_ADK_FILES: list[str] = []
+
 
 @dataclass
 class TemplateConfig:
@@ -77,12 +80,10 @@ def get_available_agents(deployment_target: str | None = None) -> dict:
         deployment_target: Optional deployment target to filter agents
     """
     # Define priority agents that should appear first
-    PRIORITY_AGENTS = [
-        "langgraph_base_react"  # Add other priority agents here as needed
-    ]
+    PRIORITY_AGENTS = ["adk_base", "agentic_rag", "langgraph_base_react"]
 
     agents_list = []
-    priority_agents = []
+    priority_agents_dict = dict.fromkeys(PRIORITY_AGENTS)  # Track priority agents
     agents_dir = pathlib.Path(__file__).parent.parent.parent.parent / "agents"
 
     for agent_dir in agents_dir.iterdir():
@@ -109,16 +110,21 @@ def get_available_agents(deployment_target: str | None = None) -> dict:
 
                     # Add to priority list or regular list based on agent name
                     if agent_name in PRIORITY_AGENTS:
-                        priority_agents.append(agent_info)
+                        priority_agents_dict[agent_name] = agent_info
                     else:
                         agents_list.append(agent_info)
                 except Exception as e:
                     logging.warning(f"Could not load agent from {agent_dir}: {e}")
 
-    # Only sort the non-priority agents
+    # Sort the non-priority agents
     agents_list.sort(key=lambda x: x["name"])
 
-    # Combine priority agents with regular agents (no sorting of priority_agents)
+    # Create priority agents list in the exact order specified
+    priority_agents = [
+        info for name, info in priority_agents_dict.items() if info is not None
+    ]
+
+    # Combine priority agents with regular agents
     combined_agents = priority_agents + agents_list
 
     # Convert to numbered dictionary starting from 1
@@ -497,29 +503,27 @@ def process_template(
             extra_deps = template_config.get("settings", {}).get(
                 "extra_dependencies", []
             )
-            otel_instrumentations = get_otel_instrumentations(dependencies=extra_deps)
-
             # Get frontend type from template config
             frontend_type = template_config.get("settings", {}).get(
                 "frontend_type", DEFAULT_FRONTEND
             )
-
+            tags = template_config.get("settings", {}).get("tags", ["None"])
             cookiecutter_config = {
                 "project_name": "my-project",
                 "agent_name": agent_name,
                 "package_version": get_current_version(),
                 "agent_description": template_config.get("description", ""),
+                "tags": tags,
                 "deployment_target": deployment_target or "",
                 "frontend_type": frontend_type,
                 "extra_dependencies": [extra_deps],
-                "otel_instrumentations": otel_instrumentations,
                 "data_ingestion": include_data_ingestion,  # Use explicit flag for cookiecutter
                 "datastore_type": datastore if datastore else "",
                 "_copy_without_render": [
                     "*.ipynb",  # Don't render notebooks
                     "*.json",  # Don't render JSON files
                     "frontend/*",  # Don't render frontend directory
-                    "tests/*",  # Don't render tests directory
+                    # "tests/*",  # Don't render tests directory
                     "notebooks/*",  # Don't render notebooks directory
                     ".git/*",  # Don't render git directory
                     "__pycache__/*",  # Don't render cache
@@ -566,6 +570,17 @@ def process_template(
                 shutil.copytree(output_dir, final_destination, dirs_exist_ok=True)
                 logging.debug(f"Project successfully created at {final_destination}")
 
+                # Delete appropriate files based on ADK tag
+                if "adk" in tags:
+                    files_to_delete = [final_destination / f for f in NON_ADK_FILES]
+                else:
+                    files_to_delete = [final_destination / f for f in ADK_FILES]
+
+                for file_path in files_to_delete:
+                    if file_path.exists():
+                        file_path.unlink()
+                        logging.debug(f"Deleted {file_path}")
+
                 # After copying template files, handle the lock file
                 if deployment_target:
                     # Get the source lock file path
@@ -675,6 +690,11 @@ def copy_files(
 
 def copy_frontend_files(frontend_type: str, project_template: pathlib.Path) -> None:
     """Copy files from the specified frontend folder directly to project root."""
+    # Skip copying if frontend_type is "None"
+    if frontend_type == "None":
+        logging.debug("Frontend type is 'None', skipping frontend files")
+        return
+
     # Use default frontend if none specified
     frontend_type = frontend_type or DEFAULT_FRONTEND
 
@@ -715,16 +735,3 @@ def copy_deployment_files(
         )
     else:
         logging.warning(f"Deployment target directory not found: {deployment_path}")
-
-
-def get_otel_instrumentations(dependencies: list) -> list[list[str]]:
-    """Returns OpenTelemetry instrumentation statements for enabled dependencies."""
-    otel_deps = {
-        "langgraph": "Instruments.LANGCHAIN",
-        "crewai": "Instruments.CREW",
-    }
-    imports = []
-    for dep in dependencies:
-        if any(otel_dep in dep for otel_dep in otel_deps):
-            imports.append(otel_deps[next(key for key in otel_deps if key in dep)])
-    return [imports]

src/data_ingestion/data_ingestion_pipeline/components/ingest_data.py

@@ -109,7 +109,8 @@ def ingest_data(
         operation = schema_client.update_schema(
             request=discoveryengine.UpdateSchemaRequest(
                 schema=new_schema, allow_missing=True
-            )
+            ),
+            timeout=1800,
         )
         logging.info(f"Waiting for schema update operation: {operation.operation.name}")
         operation.result()

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
@@ -66,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
@@ -107,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,6 +165,11 @@ class AgentEngineApp:
         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.
 
@@ -140,6 +186,7 @@ class AgentEngineApp:
             "": ["query", "register_feedback"],
             "stream": ["stream_query"],
         }
+{%- endif %}
 
 
 def deploy_agent_engine_app(
@@ -162,14 +209,18 @@ 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 = {
-        "agent_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}")

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 %}