chore: import upstream snapshot with attribution
Continuous Integration / Pre-commit Linter (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Copybara PR Handler / close-imported-pr (push) Waiting to run
Continuous Integration / Pre-commit Linter (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Copybara PR Handler / close-imported-pr (push) Waiting to run
This commit is contained in:
@@ -0,0 +1,13 @@
|
||||
# Copyright 2026 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.
|
||||
@@ -0,0 +1,87 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Microbenchmark for building LLM request contents from session history.
|
||||
|
||||
Times contents._get_contents over a large history with sizeable function-call
|
||||
payloads. To run the benchmark:
|
||||
|
||||
uv run python tests/benchmarks/benchmark_contents.py
|
||||
"""
|
||||
|
||||
from google.adk.events.event import Event
|
||||
from google.adk.flows.llm_flows import contents
|
||||
from google.genai import types
|
||||
import google_benchmark
|
||||
|
||||
_NUM_EVENTS = 500
|
||||
_PAYLOAD_SIZE = 200
|
||||
|
||||
|
||||
def _make_events(num_events: int, payload_size: int) -> list[Event]:
|
||||
"""Builds a session history with sizeable function-call payloads."""
|
||||
payload = {f"k{i}": list(range(payload_size)) for i in range(10)}
|
||||
events = [
|
||||
Event(
|
||||
invocation_id="inv0",
|
||||
author="user",
|
||||
content=types.UserContent("start"),
|
||||
)
|
||||
]
|
||||
for i in range(num_events):
|
||||
call_id = f"adk-call-{i}"
|
||||
events.append(
|
||||
Event(
|
||||
invocation_id=f"inv{i}",
|
||||
author="agent",
|
||||
content=types.Content(
|
||||
role="model",
|
||||
parts=[
|
||||
types.Part(
|
||||
function_call=types.FunctionCall(
|
||||
id=call_id, name="tool", args=dict(payload)
|
||||
)
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
events.append(
|
||||
Event(
|
||||
invocation_id=f"inv{i}",
|
||||
author="agent",
|
||||
content=types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
types.Part(
|
||||
function_response=types.FunctionResponse(
|
||||
id=call_id, name="tool", response=dict(payload)
|
||||
)
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
return events
|
||||
|
||||
|
||||
@google_benchmark.register
|
||||
def get_contents(state):
|
||||
events = _make_events(_NUM_EVENTS, _PAYLOAD_SIZE)
|
||||
while state:
|
||||
contents._get_contents(None, events, "agent")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
google_benchmark.main()
|
||||
@@ -0,0 +1,10 @@
|
||||
# Copy as .env file and fill your values below to run integration tests.
|
||||
|
||||
# Choose Backend: GOOGLE_AI_ONLY | VERTEX_ONLY | BOTH (default)
|
||||
TEST_BACKEND=BOTH
|
||||
|
||||
# ML Dev backend config
|
||||
GOOGLE_API_KEY=YOUR_VALUE_HERE
|
||||
# Vertex backend config
|
||||
GOOGLE_CLOUD_PROJECT=YOUR_VALUE_HERE
|
||||
GOOGLE_CLOUD_LOCATION=YOUR_VALUE_HERE
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright 2026 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 pytest
|
||||
|
||||
# This allows pytest to show the values of the asserts.
|
||||
pytest.register_assert_rewrite('tests.integration.utils')
|
||||
@@ -0,0 +1,119 @@
|
||||
# Copyright 2026 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 logging
|
||||
import os
|
||||
from typing import Literal
|
||||
import warnings
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from google.adk import Agent
|
||||
from pytest import fixture
|
||||
from pytest import FixtureRequest
|
||||
from pytest import hookimpl
|
||||
from pytest import Metafunc
|
||||
|
||||
from .utils import TestRunner
|
||||
|
||||
logger = logging.getLogger('google_adk.' + __name__)
|
||||
|
||||
|
||||
def load_env_for_tests():
|
||||
dotenv_path = os.path.join(os.path.dirname(__file__), '.env')
|
||||
if not os.path.exists(dotenv_path):
|
||||
warnings.warn(
|
||||
f'Missing .env file at {dotenv_path}. See dotenv.sample for an example.'
|
||||
)
|
||||
else:
|
||||
load_dotenv(dotenv_path, override=True, verbose=True)
|
||||
if 'GOOGLE_API_KEY' not in os.environ:
|
||||
warnings.warn(
|
||||
'Missing GOOGLE_API_KEY in the environment variables. GOOGLE_AI backend'
|
||||
' integration tests will fail.'
|
||||
)
|
||||
for env_var in [
|
||||
'GOOGLE_CLOUD_PROJECT',
|
||||
'GOOGLE_CLOUD_LOCATION',
|
||||
]:
|
||||
if env_var not in os.environ:
|
||||
warnings.warn(
|
||||
f'Missing {env_var} in the environment variables. Vertex backend'
|
||||
' integration tests will fail.'
|
||||
)
|
||||
|
||||
|
||||
load_env_for_tests()
|
||||
|
||||
BackendType = Literal['GOOGLE_AI', 'VERTEX']
|
||||
|
||||
|
||||
@fixture
|
||||
def agent_runner(request: FixtureRequest) -> TestRunner:
|
||||
assert isinstance(request.param, dict)
|
||||
|
||||
if 'agent' in request.param:
|
||||
assert isinstance(request.param['agent'], Agent)
|
||||
return TestRunner(request.param['agent'])
|
||||
elif 'agent_name' in request.param:
|
||||
assert isinstance(request.param['agent_name'], str)
|
||||
return TestRunner.from_agent_name(request.param['agent_name'])
|
||||
|
||||
raise NotImplementedError('Must provide agent or agent_name.')
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def llm_backend(request: FixtureRequest):
|
||||
# Set backend environment value.
|
||||
original_val = os.environ.get('GOOGLE_GENAI_USE_ENTERPRISE')
|
||||
backend_type = request.param
|
||||
if backend_type == 'GOOGLE_AI':
|
||||
os.environ['GOOGLE_GENAI_USE_ENTERPRISE'] = '0'
|
||||
else:
|
||||
os.environ['GOOGLE_GENAI_USE_ENTERPRISE'] = '1'
|
||||
|
||||
yield # Run the test
|
||||
|
||||
# Restore the environment
|
||||
if original_val is None:
|
||||
os.environ.pop('GOOGLE_GENAI_USE_ENTERPRISE', None)
|
||||
else:
|
||||
os.environ['GOOGLE_GENAI_USE_ENTERPRISE'] = original_val
|
||||
|
||||
|
||||
@hookimpl(tryfirst=True)
|
||||
def pytest_generate_tests(metafunc: Metafunc):
|
||||
if llm_backend.__name__ in metafunc.fixturenames:
|
||||
if not _is_explicitly_marked(llm_backend.__name__, metafunc):
|
||||
test_backend = os.environ.get('TEST_BACKEND', 'BOTH')
|
||||
if test_backend == 'GOOGLE_AI_ONLY':
|
||||
metafunc.parametrize(llm_backend.__name__, ['GOOGLE_AI'], indirect=True)
|
||||
elif test_backend == 'VERTEX_ONLY':
|
||||
metafunc.parametrize(llm_backend.__name__, ['VERTEX'], indirect=True)
|
||||
elif test_backend == 'BOTH':
|
||||
metafunc.parametrize(
|
||||
llm_backend.__name__, ['GOOGLE_AI', 'VERTEX'], indirect=True
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f'Invalid TEST_BACKEND value: {test_backend}, should be one of'
|
||||
' [GOOGLE_AI_ONLY, VERTEX_ONLY, BOTH]'
|
||||
)
|
||||
|
||||
|
||||
def _is_explicitly_marked(mark_name: str, metafunc: Metafunc) -> bool:
|
||||
if hasattr(metafunc.function, 'pytestmark'):
|
||||
for mark in metafunc.function.pytestmark:
|
||||
if mark.name == 'parametrize' and mark_name in mark.args[0]:
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright 2026 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.
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,88 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk import Agent
|
||||
from google.genai import types
|
||||
|
||||
new_message = types.Content(
|
||||
role="user",
|
||||
parts=[types.Part.from_text(text="Count a number")],
|
||||
)
|
||||
|
||||
google_agent_1 = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="agent_1",
|
||||
description="The first agent in the team.",
|
||||
instruction="Just say 1",
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
),
|
||||
)
|
||||
|
||||
google_agent_2 = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="agent_2",
|
||||
description="The second agent in the team.",
|
||||
instruction="Just say 2",
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.2,
|
||||
safety_settings=[{
|
||||
"category": "HARM_CATEGORY_HATE_SPEECH",
|
||||
"threshold": "BLOCK_ONLY_HIGH",
|
||||
}],
|
||||
),
|
||||
)
|
||||
|
||||
google_agent_3 = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="agent_3",
|
||||
description="The third agent in the team.",
|
||||
instruction="Just say 3",
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.5,
|
||||
safety_settings=[{
|
||||
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
"threshold": "BLOCK_NONE",
|
||||
}],
|
||||
),
|
||||
)
|
||||
|
||||
google_agent_with_instruction_in_config = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="agent",
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.5, system_instruction="Count 1"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def function():
|
||||
pass
|
||||
|
||||
|
||||
google_agent_with_tools_in_config = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="agent",
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.5, tools=[function]
|
||||
),
|
||||
)
|
||||
|
||||
google_agent_with_response_schema_in_config = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="agent",
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.5, response_schema={"key": "value"}
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,67 @@
|
||||
# Instructions
|
||||
|
||||
## Run Evaluation
|
||||
|
||||
1. Set environment variables in your terminal:
|
||||
|
||||
```shell
|
||||
export GOOGLE_GENAI_USE_ENTERPRISE=FALSE
|
||||
export GOOGLE_API_KEY=<your_api_key>
|
||||
export GOOGLE_CLOUD_PROJECT=<your_bigquery_project>
|
||||
```
|
||||
1. Change to the current directory:
|
||||
|
||||
```shell
|
||||
cd tests/integration/fixture/bigquery_agent/
|
||||
```
|
||||
1. Customize the evaluation dataset to the environment `GOOGLE_CLOUD_PROJECT`
|
||||
by replacing the placeholder to the real project set in your environment:
|
||||
|
||||
```shell
|
||||
sed -e "s:\${GOOGLE_CLOUD_PROJECT}:${GOOGLE_CLOUD_PROJECT}:g" simple.test.json -i
|
||||
```
|
||||
1. Run the following command as per https://google.github.io/adk-docs/evaluate/#3-adk-eval-run-evaluations-via-the-cli:
|
||||
|
||||
```shell
|
||||
adk eval . simple.test.json --config_file_path=test_config.json
|
||||
```
|
||||
|
||||
If it fails, re-run with `--print_detailed_results` flag to see more details
|
||||
on turn-by-turn evaluation.
|
||||
|
||||
## Generate Evaluation dataset
|
||||
|
||||
1. Set environment variables in your terminal:
|
||||
|
||||
```shell
|
||||
export GOOGLE_GENAI_USE_ENTERPRISE=FALSE
|
||||
export GOOGLE_API_KEY=<your_api_key>
|
||||
export GOOGLE_CLOUD_PROJECT=<your_bigquery_project>
|
||||
```
|
||||
1. Set up google [application default credentials](https://cloud.google.com/docs/authentication/provide-credentials-adc)
|
||||
on your machine.
|
||||
|
||||
```shell
|
||||
gcloud auth application-default login
|
||||
```
|
||||
1. Change to the directory containing agent folder:
|
||||
|
||||
```shell
|
||||
cd tests/integration/fixture/
|
||||
```
|
||||
1. Run the following command to start the ADK web app:
|
||||
|
||||
```shell
|
||||
adk web
|
||||
```
|
||||
1. Open the ADK web UI in your browser http://127.0.0.1:8000/dev-ui/?app=bigquery_agent.
|
||||
1. Create an evaluation dataset by following [these steps](https://google.github.io/adk-docs/evaluate/#1-adk-web-run-evaluations-via-the-web-ui).
|
||||
This would generate file `bigquery_agent/simple.evalset.json`.
|
||||
1. Note that this evaluation data would be tied to the agent interaction in the
|
||||
`GOOGLE_CLOUD_PROJECT` set in your environment. To normalize it by replacing
|
||||
the real project set in your environment to a placeholder, let's run the
|
||||
following command:
|
||||
|
||||
```shell
|
||||
sed -e "s:${GOOGLE_CLOUD_PROJECT}:\${GOOGLE_CLOUD_PROJECT}:g" bigquery_agent/simple.evalset.json > bigquery_agent/simple.test.json
|
||||
```
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,75 @@
|
||||
# Copyright 2026 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 os
|
||||
|
||||
from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
|
||||
from google.adk.tools.bigquery.bigquery_toolset import BigQueryToolset
|
||||
from google.adk.tools.bigquery.config import BigQueryToolConfig
|
||||
from google.adk.tools.bigquery.config import WriteMode
|
||||
import google.auth
|
||||
|
||||
# Check necessary environment variables
|
||||
if not (google_cloud_project_id := os.getenv("GOOGLE_CLOUD_PROJECT")):
|
||||
raise ValueError(
|
||||
"GOOGLE_CLOUD_PROJECT environment variable is not set. Please set it"
|
||||
" to the GCP project ID where your BigQuery jobs would be run."
|
||||
)
|
||||
|
||||
# Define an appropriate application name
|
||||
BIGQUERY_AGENT_NAME = "adk_eval_bigquery_agent"
|
||||
|
||||
|
||||
# Define BigQuery tool config with write mode set to allowed. Note that this is
|
||||
# only to demonstrate the full capability of the BigQuery tools. In production
|
||||
# you may want to change to BLOCKED (default write mode, effectively makes the
|
||||
# tool read-only) or PROTECTED (only allows writes in the anonymous dataset of a
|
||||
# BigQuery session) write mode.
|
||||
tool_config = BigQueryToolConfig(
|
||||
write_mode=WriteMode.BLOCKED,
|
||||
application_name=BIGQUERY_AGENT_NAME,
|
||||
compute_project_id=google_cloud_project_id,
|
||||
)
|
||||
|
||||
# Initialize the tools to use the application default credentials.
|
||||
# https://cloud.google.com/docs/authentication/provide-credentials-adc
|
||||
application_default_credentials, _ = google.auth.default()
|
||||
credentials_config = BigQueryCredentialsConfig(
|
||||
credentials=application_default_credentials
|
||||
)
|
||||
|
||||
bigquery_toolset = BigQueryToolset(
|
||||
credentials_config=credentials_config, bigquery_tool_config=tool_config
|
||||
)
|
||||
|
||||
# The variable name `root_agent` determines what your root agent is for the
|
||||
# debug CLI
|
||||
root_agent = LlmAgent(
|
||||
model="gemini-2.5-flash",
|
||||
name=BIGQUERY_AGENT_NAME,
|
||||
description=(
|
||||
"Agent to answer questions about BigQuery data and models and execute"
|
||||
" SQL queries."
|
||||
),
|
||||
instruction=f"""\
|
||||
You are a data science agent with access to several BigQuery tools.
|
||||
Make use of those tools to answer the user's questions.
|
||||
|
||||
You must use the project id {google_cloud_project_id} for running SQL
|
||||
queries and generating data insights
|
||||
""",
|
||||
tools=[bigquery_toolset],
|
||||
)
|
||||
@@ -0,0 +1,538 @@
|
||||
{
|
||||
"eval_set_id": "simple",
|
||||
"name": "simple",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "penguins exploration",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "e-81734a2f-60dc-4c7e-9b05-29a2c18b7651",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "Hi, what can you do?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "I can help you with BigQuery. I can:\n* List datasets and tables in a project\n* Get information about datasets and tables\n* Execute SQL queries\n* Forecast time series data\n* Answer questions about data in BigQuery tables using natural language.\n\nWhat would you like to do?\n"
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757645225.080524
|
||||
},
|
||||
{
|
||||
"invocation_id": "e-78026d6b-f3d5-4807-8122-8872efb024d6",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "which tools do you have access to?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": "CqsCAdHtim-ZZuTPQdEgbCYXWniB28PuU35grogIwz7a2X0PG9eopcLAT0hXOv90Zu5b9Q4iREsIAcV3znUrCjwMwRrW_G3QYH7jfaK5oEvunfD-yOUj-V9wjc_hAyKon4ogXs7nxX9v_sAfz1uh48cmiMBcNcJd1dkiBbCufChWQbHFsVghXAnetSRm21H2rRgPuxvpf4_mWXzxgYZddOuf6bYmUI3kEcwQDqbocOk2u-ghG44KnFAXOhqBBu8eUeJM8m7uUWevEtxXIclJ14crWCjWADzAof80VX_rLucQ5sPE5wfTvlFIDECuGapnxmxAVB8hNCw9iwlJEhNUVVSbPPfvjwDSq9s99w-Rbu8yLDA5YwSru-q1qIrxbXTWuNlh9fwCdKDQAgiXUo0=",
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "I have access to the following tools:\n\n* `get_dataset_info(project_id, dataset_id)`: Get metadata information about a BigQuery dataset.\n* `get_table_info(project_id, dataset_id, table_id)`: Get metadata information about a BigQuery table.\n* `list_dataset_ids(project_id)`: List BigQuery dataset IDs in a Google Cloud project.\n* `list_table_ids(project_id, dataset_id)`: List table IDs in a BigQuery dataset.\n* `execute_sql(project_id, query)`: Run a BigQuery or BigQuery ML SQL query in the project and return the result.\n* `forecast(project_id, history_data, timestamp_col, data_col, horizon, id_cols=None)`: Run a BigQuery AI time series forecast using AI.FORECAST.\n* `ask_data_insights(project_id, user_query_with_context, table_references)`: Answers questions about structured data in BigQuery tables using natural language."
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757645280.602244
|
||||
},
|
||||
{
|
||||
"invocation_id": "e-47aa5d37-6bba-4a2d-8eae-e79f47f347c4",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "Are there any ML datasets in the bigquery public data?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": "Cq8DAdHtim9T93wEXUnP9hCJ0SGm79kvsT5VJBWIL1xr8Z0TBLYIQSVqogdU3mB3XkwHkqKT7hBGBY11yuHwfcohBakiOco71gRXOlhf0XGlNZNIUUObnOdY2swLsmpJDbOtRLxgU9OZ0JhHlC0fkrQj9Ab2wt5A8VFkuBUQaEB-XcJYes8Zo0TfU79nrlKrfIINPHsXEuBdq4biDipPss57EZwgs8HiwdeUCXeMwcYS1NFquYUmFLqnsAf9Xik5k3yEx9iQyF7VtOPNGp9sKC2zh7Euz5bpgiHjRTV41hw5QWQk8Q38hKOS7G2jLLXPO8v63sn9LgIzCHNJUd9j-IU8v5gtgU9CNfG4o7icU7GD77KrJx6etaxHQiwSfMjPR_6tZ_ft4-eIh7kRQdJo-GrtyDJJBNV-5s61G0qMGqbS7JY4RfH5_cT8UUSEamkU1eJGN7pZIAKc7FXaCHlGUOQClx-9D7XXQMDd2gVk1pyQvMeruzmykywSGPffnDZJQ9kVwW9pl6urnk8hX4ZGnXM_DTjaFZFF_Gb1dhpr0Uuy0LWk1IqcfbS95YJIFipFXPo=",
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "Yes, there are a few ML-related datasets in the `bigquery-public-data` project:\n\n* `aml_ai_input_dataset`\n* `bigqueryml_ncaa`\n* `ml_datasets`\n* `ml_datasets_uscentral1`"
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": "adk-4d524ad4-6a3b-49dd-a44e-763f956f89c1",
|
||||
"args": {
|
||||
"project_id": "bigquery-public-data"
|
||||
},
|
||||
"name": "list_dataset_ids"
|
||||
}
|
||||
],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757645323.568069
|
||||
},
|
||||
{
|
||||
"invocation_id": "e-76115967-fe7b-4aee-92c8-39d3fd1b35c0",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "When was ml_datasets created and which location it is in?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "The `ml_datasets` dataset was created on `1553208775542` (which is March 21, 2019, 19:12:55 UTC) and is located in the `US` region."
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": "adk-2cba4269-dde6-4b62-a38a-1004f05452a0",
|
||||
"args": {
|
||||
"dataset_id": "ml_datasets",
|
||||
"project_id": "bigquery-public-data"
|
||||
},
|
||||
"name": "get_dataset_info"
|
||||
}
|
||||
],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757645409.36703
|
||||
},
|
||||
{
|
||||
"invocation_id": "e-9a5e0a24-f7cb-4e50-9a61-430a11837e82",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "alright, which tables does it have?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "The `ml_datasets` dataset contains the following tables:\n\n* `census_adult_income`\n* `census_adult_income_kmeans_predictions`\n* `credit_card_default`\n* `holidays_and_events_for_forecasting`\n* `iris`\n* `penguins`\n* `ulb_fraud_detection`"
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": "adk-089bf143-422f-4516-b6d9-810773b03de7",
|
||||
"args": {
|
||||
"dataset_id": "ml_datasets",
|
||||
"project_id": "bigquery-public-data"
|
||||
},
|
||||
"name": "list_table_ids"
|
||||
}
|
||||
],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757645478.923206
|
||||
},
|
||||
{
|
||||
"invocation_id": "e-05b241ed-9678-49dc-8ead-25baab6d4002",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "tell me more details about the penguins table"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "The `penguins` table was created on `1619804743188` (which is April 30, 2021, 00:25:43 UTC) and contains 344 rows. It has the following schema:\n\n* `species`: STRING (REQUIRED)\n* `island`: STRING (NULLABLE)\n* `culmen_length_mm`: FLOAT (NULLABLE)\n* `culmen_depth_mm`: FLOAT (NULLABLE)\n* `flipper_length_mm`: FLOAT (NULLABLE)\n* `body_mass_g`: FLOAT (NULLABLE)\n* `sex`: STRING (NULLABLE)"
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": "adk-8893d916-9445-4458-8998-ea3e283d534d",
|
||||
"args": {
|
||||
"dataset_id": "ml_datasets",
|
||||
"project_id": "bigquery-public-data",
|
||||
"table_id": "penguins"
|
||||
},
|
||||
"name": "get_table_info"
|
||||
}
|
||||
],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757645526.958228
|
||||
},
|
||||
{
|
||||
"invocation_id": "e-d7b85acb-b7f3-4b61-b66d-bba77c01e13e",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "can you tell me population of penguins per island?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "Here is the population of penguins per island:\n\n* **Dream:** 124\n* **Biscoe:** 168\n* **Torgersen:** 52"
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": "adk-33ff1809-bbfb-45e7-a07c-25fc9c932260",
|
||||
"args": {
|
||||
"query": "SELECT island, COUNT(*) AS population FROM bigquery-public-data.ml_datasets.penguins GROUP BY island",
|
||||
"project_id": "${GOOGLE_CLOUD_PROJECT}"
|
||||
},
|
||||
"name": "execute_sql"
|
||||
}
|
||||
],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757645563.777316
|
||||
},
|
||||
{
|
||||
"invocation_id": "e-3b46dd99-4d25-45e5-b16c-2f055de56f47",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "can you get me insights about any correlation between body mass and flipper length of penguins?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "Based on the data, there appears to be a positive correlation between the body mass and flipper length of penguins. Generally, as the flipper length in millimeters increases, the body mass in grams also tends to increase."
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": "adk-266b6452-50b8-4d3b-ae23-d09effa41e0b",
|
||||
"args": {
|
||||
"project_id": "${GOOGLE_CLOUD_PROJECT}",
|
||||
"user_query_with_context": "Is there any correlation between body mass and flipper length of penguins?",
|
||||
"table_references": [
|
||||
{
|
||||
"datasetId": "ml_datasets",
|
||||
"tableId": "penguins",
|
||||
"projectId": "bigquery-public-data"
|
||||
}
|
||||
]
|
||||
},
|
||||
"name": "ask_data_insights"
|
||||
}
|
||||
],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757645706.447186
|
||||
},
|
||||
{
|
||||
"invocation_id": "e-357fb858-adae-4a7c-80fc-fcadf52a20fd",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "how do the penguins tables in ml_datasets and ml_datasets_uscentral compare?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "Here's a comparison of the `penguins` tables in `ml_datasets` and `ml_datasets_uscentral1`:\n\n**`bigquery-public-data:ml_datasets.penguins`**\n\n* **Creation Time:** 1619804743188 (April 30, 2021, 00:25:43 UTC)\n* **Location:** US\n* **Number of Rows:** 344\n* **Schema:**\n * `species`: STRING (REQUIRED)\n * `island`: STRING (NULLABLE)\n * `culmen_length_mm`: FLOAT (NULLABLE)\n * `culmen_depth_mm`: FLOAT (NULLABLE)\n * `flipper_length_mm`: FLOAT (NULLABLE)\n * `body_mass_g`: FLOAT (NULLABLE)\n * `sex`: STRING (NULLABLE)\n\n**`bigquery-public-data:ml_datasets_uscentral1.penguins`**\n\n* **Creation Time:** 1634583926374 (October 18, 2021, 23:05:26 UTC)\n* **Location:** us-central1\n* **Number of Rows:** 344\n* **Schema:** (Identical to `ml_datasets.penguins`)\n * `species`: STRING (REQUIRED)\n * `island`: STRING (NULLABLE)\n * `culmen_length_mm`: FLOAT (NULLABLE)\n * `culmen_depth_mm`: FLOAT (NULLABLE)\n * `flipper_length_mm`: FLOAT (NULLABLE)\n * `body_mass_g`: FLOAT (NULLABLE)\n * `sex`: STRING (NULLABLE)\n\n**Key Differences:**\n\n* **Creation Time:** The table in `ml_datasets` was created earlier (April 2021) than the one in `ml_datasets_uscentral1` (October 2021).\n* **Location:** The `ml_datasets.penguins` table is in the `US` multi-region, while `ml_datasets_uscentral1.penguins` is specifically in the `us-central1` region.\n\n**Similarities:**\n\n* Both tables have the **same number of rows (344)**.\n* Both tables have an **identical schema**, with the same column names, types, and modes.\n\nIt appears that `bigquery-public-data:ml_datasets_uscentral1.penguins` is likely a regional copy or a later version of the `bigquery-public-data:ml_datasets.penguins` table."
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": "adk-8ff49a41-b42a-4f2b-8408-bda00712791b",
|
||||
"args": {
|
||||
"project_id": "bigquery-public-data",
|
||||
"dataset_id": "ml_datasets",
|
||||
"table_id": "penguins"
|
||||
},
|
||||
"name": "get_table_info"
|
||||
},
|
||||
{
|
||||
"id": "adk-c0e9a831-300d-4805-a127-6df1d7afc712",
|
||||
"args": {
|
||||
"project_id": "bigquery-public-data",
|
||||
"dataset_id": "ml_datasets_uscentral1",
|
||||
"table_id": "penguins"
|
||||
},
|
||||
"name": "get_table_info"
|
||||
}
|
||||
],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757646232.887191
|
||||
},
|
||||
{
|
||||
"invocation_id": "e-72ddf0dd-4fee-41dd-adb6-26cd98baeb25",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "can you tell the population per island from the uscentral1 table?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "Here is the population of penguins per island from the `ml_datasets_uscentral1` table:\n\n* **Dream:** 124\n* **Biscoe:** 168\n* **Torgersen:** 52"
|
||||
}
|
||||
],
|
||||
"role": null
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": "adk-febf7cf1-3465-4393-9882-ef860502238f",
|
||||
"args": {
|
||||
"project_id": "${GOOGLE_CLOUD_PROJECT}",
|
||||
"query": "SELECT island, COUNT(*) AS population FROM `bigquery-public-data`.`ml_datasets_uscentral1`.`penguins` GROUP BY island"
|
||||
},
|
||||
"name": "execute_sql"
|
||||
}
|
||||
],
|
||||
"tool_responses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1757646339.73431
|
||||
}
|
||||
],
|
||||
"session_input": {
|
||||
"app_name": "bigquery_agent",
|
||||
"user_id": "user",
|
||||
"state": {}
|
||||
},
|
||||
"creation_timestamp": 1757648114.6285758
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1757648101.7927744
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"criteria": {
|
||||
"tool_trajectory_avg_score": 0.7,
|
||||
"response_match_score": 0.7
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,105 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from google.adk.agents.callback_context import CallbackContext
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.llm_agent import Agent
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.models.llm_response import LlmResponse
|
||||
from google.genai import types
|
||||
|
||||
|
||||
def before_agent_call_end_invocation(
|
||||
callback_context: CallbackContext,
|
||||
) -> types.Content:
|
||||
return types.Content(
|
||||
role='model',
|
||||
parts=[types.Part(text='End invocation event before agent call.')],
|
||||
)
|
||||
|
||||
|
||||
def before_agent_call(
|
||||
invocation_context: InvocationContext,
|
||||
) -> types.Content:
|
||||
return types.Content(
|
||||
role='model',
|
||||
parts=[types.Part.from_text(text='Plain text event before agent call.')],
|
||||
)
|
||||
|
||||
|
||||
def before_model_call_end_invocation(
|
||||
callback_context: CallbackContext, llm_request: LlmRequest
|
||||
) -> LlmResponse:
|
||||
return LlmResponse(
|
||||
content=types.Content(
|
||||
role='model',
|
||||
parts=[
|
||||
types.Part.from_text(
|
||||
text='End invocation event before model call.'
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def before_model_call(
|
||||
invocation_context: InvocationContext, request: LlmRequest
|
||||
) -> LlmResponse:
|
||||
request.config.system_instruction = 'Just return 999 as response.'
|
||||
return LlmResponse(
|
||||
content=types.Content(
|
||||
role='model',
|
||||
parts=[
|
||||
types.Part.from_text(
|
||||
text='Update request event before model call.'
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def after_model_call(
|
||||
callback_context: CallbackContext,
|
||||
llm_response: LlmResponse,
|
||||
) -> Optional[LlmResponse]:
|
||||
content = llm_response.content
|
||||
if not content or not content.parts or not content.parts[0].text:
|
||||
return
|
||||
|
||||
content.parts[0].text += 'Update response event after model call.'
|
||||
return llm_response
|
||||
|
||||
|
||||
before_agent_callback_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='before_agent_callback_agent',
|
||||
instruction='echo 1',
|
||||
before_agent_callback=before_agent_call_end_invocation,
|
||||
)
|
||||
|
||||
before_model_callback_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='before_model_callback_agent',
|
||||
instruction='echo 2',
|
||||
before_model_callback=before_model_call_end_invocation,
|
||||
)
|
||||
|
||||
after_model_callback_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='after_model_callback_agent',
|
||||
instruction='Say hello',
|
||||
after_model_callback=after_model_call,
|
||||
)
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from typing import List
|
||||
from typing import Union
|
||||
|
||||
from google.adk import Agent
|
||||
from google.adk.tools.tool_context import ToolContext
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
def update_fc(
|
||||
data_one: str,
|
||||
data_two: Union[int, float, str],
|
||||
data_three: list[str],
|
||||
data_four: List[Union[int, float, str]],
|
||||
tool_context: ToolContext,
|
||||
):
|
||||
"""Simply ask to update these variables in the context"""
|
||||
tool_context.actions.update_state("data_one", data_one)
|
||||
tool_context.actions.update_state("data_two", data_two)
|
||||
tool_context.actions.update_state("data_three", data_three)
|
||||
tool_context.actions.update_state("data_four", data_four)
|
||||
|
||||
|
||||
root_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="root_agent",
|
||||
instruction="Call tools",
|
||||
flow="auto",
|
||||
tools=[update_fc],
|
||||
)
|
||||
@@ -0,0 +1,582 @@
|
||||
{
|
||||
"id": "ead43200-b575-4241-9248-233b4be4f29a",
|
||||
"context": {
|
||||
"_time": "2024-12-01 09:02:43.531503",
|
||||
"data_one": "RRRR",
|
||||
"data_two": "3.141529",
|
||||
"data_three": [
|
||||
"apple",
|
||||
"banana"
|
||||
],
|
||||
"data_four": [
|
||||
"1",
|
||||
"hello",
|
||||
"3.14"
|
||||
]
|
||||
},
|
||||
"events": [
|
||||
{
|
||||
"invocation_id": "6BGrtKJu",
|
||||
"author": "user",
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "hi"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"options": {},
|
||||
"id": "ltzQTqR4",
|
||||
"timestamp": 1733043686.8428597
|
||||
},
|
||||
{
|
||||
"invocation_id": "6BGrtKJu",
|
||||
"author": "root_agent",
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "Hello! 👋 How can I help you today? \n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"options": {
|
||||
"partial": false
|
||||
},
|
||||
"id": "ClSROx8b",
|
||||
"timestamp": 1733043688.1030986
|
||||
},
|
||||
{
|
||||
"invocation_id": "M3dUcVa8",
|
||||
"author": "user",
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "update data_one to be RRRR, data_two to be 3.141529, data_three to be apple and banana, data_four to be 1, hello, and 3.14"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"options": {},
|
||||
"id": "yxigGwIZ",
|
||||
"timestamp": 1733043745.9900541
|
||||
},
|
||||
{
|
||||
"invocation_id": "M3dUcVa8",
|
||||
"author": "root_agent",
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"function_call": {
|
||||
"args": {
|
||||
"data_four": [
|
||||
"1",
|
||||
"hello",
|
||||
"3.14"
|
||||
],
|
||||
"data_two": "3.141529",
|
||||
"data_three": [
|
||||
"apple",
|
||||
"banana"
|
||||
],
|
||||
"data_one": "RRRR"
|
||||
},
|
||||
"name": "update_fc"
|
||||
}
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"options": {
|
||||
"partial": false
|
||||
},
|
||||
"id": "8V6de8th",
|
||||
"timestamp": 1733043747.4545543
|
||||
},
|
||||
{
|
||||
"invocation_id": "M3dUcVa8",
|
||||
"author": "root_agent",
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"function_response": {
|
||||
"name": "update_fc",
|
||||
"response": {}
|
||||
}
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"options": {
|
||||
"update_context": {
|
||||
"data_one": "RRRR",
|
||||
"data_two": "3.141529",
|
||||
"data_three": [
|
||||
"apple",
|
||||
"banana"
|
||||
],
|
||||
"data_four": [
|
||||
"1",
|
||||
"hello",
|
||||
"3.14"
|
||||
]
|
||||
},
|
||||
"function_call_event_id": "8V6de8th"
|
||||
},
|
||||
"id": "dkTj5v8B",
|
||||
"timestamp": 1733043747.457031
|
||||
},
|
||||
{
|
||||
"invocation_id": "M3dUcVa8",
|
||||
"author": "root_agent",
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "OK. I've updated the data. Anything else? \n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"options": {
|
||||
"partial": false
|
||||
},
|
||||
"id": "OZ77XR41",
|
||||
"timestamp": 1733043748.7901294
|
||||
}
|
||||
],
|
||||
"past_events": [],
|
||||
"pending_events": {},
|
||||
"artifacts": {},
|
||||
"event_logs": [
|
||||
{
|
||||
"invocation_id": "6BGrtKJu",
|
||||
"event_id": "ClSROx8b",
|
||||
"model_request": {
|
||||
"model": "gemini-2.5-flash",
|
||||
"contents": [
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "hi"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
}
|
||||
],
|
||||
"config": {
|
||||
"system_instruction": "You are an agent. Your name is root_agent.\nCall tools",
|
||||
"tools": [
|
||||
{
|
||||
"function_declarations": [
|
||||
{
|
||||
"description": "Hello",
|
||||
"name": "update_fc",
|
||||
"parameters": {
|
||||
"type": "OBJECT",
|
||||
"properties": {
|
||||
"data_one": {
|
||||
"type": "STRING"
|
||||
},
|
||||
"data_two": {
|
||||
"type": "STRING"
|
||||
},
|
||||
"data_three": {
|
||||
"type": "ARRAY",
|
||||
"items": {
|
||||
"type": "STRING"
|
||||
}
|
||||
},
|
||||
"data_four": {
|
||||
"type": "ARRAY",
|
||||
"items": {
|
||||
"any_of": [
|
||||
{
|
||||
"type": "INTEGER"
|
||||
},
|
||||
{
|
||||
"type": "NUMBER"
|
||||
},
|
||||
{
|
||||
"type": "STRING"
|
||||
}
|
||||
],
|
||||
"type": "STRING"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"model_response": {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "Hello! 👋 How can I help you today? \n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"avg_logprobs": -0.15831730915949896,
|
||||
"finish_reason": "STOP",
|
||||
"safety_ratings": [
|
||||
{
|
||||
"category": "HARM_CATEGORY_HATE_SPEECH",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.071777344,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.07080078
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.16308594,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.14160156
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_HARASSMENT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.09423828,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.037841797
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.059326172,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.02368164
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"model_version": "gemini-2.5-flash",
|
||||
"usage_metadata": {
|
||||
"candidates_token_count": 13,
|
||||
"prompt_token_count": 32,
|
||||
"total_token_count": 45
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"invocation_id": "M3dUcVa8",
|
||||
"event_id": "8V6de8th",
|
||||
"model_request": {
|
||||
"model": "gemini-2.5-flash",
|
||||
"contents": [
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "hi"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "Hello! 👋 How can I help you today? \n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "update data_one to be RRRR, data_two to be 3.141529, data_three to be apple and banana, data_four to be 1, hello, and 3.14"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
}
|
||||
],
|
||||
"config": {
|
||||
"system_instruction": "You are an agent. Your name is root_agent.\nCall tools",
|
||||
"tools": [
|
||||
{
|
||||
"function_declarations": [
|
||||
{
|
||||
"description": "Hello",
|
||||
"name": "update_fc",
|
||||
"parameters": {
|
||||
"type": "OBJECT",
|
||||
"properties": {
|
||||
"data_one": {
|
||||
"type": "STRING"
|
||||
},
|
||||
"data_two": {
|
||||
"type": "STRING"
|
||||
},
|
||||
"data_three": {
|
||||
"type": "ARRAY",
|
||||
"items": {
|
||||
"type": "STRING"
|
||||
}
|
||||
},
|
||||
"data_four": {
|
||||
"type": "ARRAY",
|
||||
"items": {
|
||||
"any_of": [
|
||||
{
|
||||
"type": "INTEGER"
|
||||
},
|
||||
{
|
||||
"type": "NUMBER"
|
||||
},
|
||||
{
|
||||
"type": "STRING"
|
||||
}
|
||||
],
|
||||
"type": "STRING"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"model_response": {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"function_call": {
|
||||
"args": {
|
||||
"data_four": [
|
||||
"1",
|
||||
"hello",
|
||||
"3.14"
|
||||
],
|
||||
"data_two": "3.141529",
|
||||
"data_three": [
|
||||
"apple",
|
||||
"banana"
|
||||
],
|
||||
"data_one": "RRRR"
|
||||
},
|
||||
"name": "update_fc"
|
||||
}
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"avg_logprobs": -2.100960955431219e-6,
|
||||
"finish_reason": "STOP",
|
||||
"safety_ratings": [
|
||||
{
|
||||
"category": "HARM_CATEGORY_HATE_SPEECH",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.12158203,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.13671875
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.421875,
|
||||
"severity": "HARM_SEVERITY_LOW",
|
||||
"severity_score": 0.24511719
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_HARASSMENT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.15722656,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.072753906
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.083984375,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.03564453
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"model_version": "gemini-2.5-flash",
|
||||
"usage_metadata": {
|
||||
"candidates_token_count": 32,
|
||||
"prompt_token_count": 94,
|
||||
"total_token_count": 126
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"invocation_id": "M3dUcVa8",
|
||||
"event_id": "OZ77XR41",
|
||||
"model_request": {
|
||||
"model": "gemini-2.5-flash",
|
||||
"contents": [
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "hi"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "Hello! 👋 How can I help you today? \n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "update data_one to be RRRR, data_two to be 3.141529, data_three to be apple and banana, data_four to be 1, hello, and 3.14"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"function_call": {
|
||||
"args": {
|
||||
"data_four": [
|
||||
"1",
|
||||
"hello",
|
||||
"3.14"
|
||||
],
|
||||
"data_two": "3.141529",
|
||||
"data_three": [
|
||||
"apple",
|
||||
"banana"
|
||||
],
|
||||
"data_one": "RRRR"
|
||||
},
|
||||
"name": "update_fc"
|
||||
}
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"function_response": {
|
||||
"name": "update_fc",
|
||||
"response": {}
|
||||
}
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
}
|
||||
],
|
||||
"config": {
|
||||
"system_instruction": "You are an agent. Your name is root_agent.\nCall tools",
|
||||
"tools": [
|
||||
{
|
||||
"function_declarations": [
|
||||
{
|
||||
"description": "Hello",
|
||||
"name": "update_fc",
|
||||
"parameters": {
|
||||
"type": "OBJECT",
|
||||
"properties": {
|
||||
"data_one": {
|
||||
"type": "STRING"
|
||||
},
|
||||
"data_two": {
|
||||
"type": "STRING"
|
||||
},
|
||||
"data_three": {
|
||||
"type": "ARRAY",
|
||||
"items": {
|
||||
"type": "STRING"
|
||||
}
|
||||
},
|
||||
"data_four": {
|
||||
"type": "ARRAY",
|
||||
"items": {
|
||||
"any_of": [
|
||||
{
|
||||
"type": "INTEGER"
|
||||
},
|
||||
{
|
||||
"type": "NUMBER"
|
||||
},
|
||||
{
|
||||
"type": "STRING"
|
||||
}
|
||||
],
|
||||
"type": "STRING"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"model_response": {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "OK. I've updated the data. Anything else? \n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"avg_logprobs": -0.22089435373033797,
|
||||
"finish_reason": "STOP",
|
||||
"safety_ratings": [
|
||||
{
|
||||
"category": "HARM_CATEGORY_HATE_SPEECH",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.04663086,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.09423828
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.18554688,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.111328125
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_HARASSMENT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.071777344,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.03112793
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probability_score": 0.043945313,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severity_score": 0.057373047
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"model_version": "gemini-2.5-flash",
|
||||
"usage_metadata": {
|
||||
"candidates_token_count": 14,
|
||||
"prompt_token_count": 129,
|
||||
"total_token_count": 143
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,115 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from typing import List
|
||||
from typing import Union
|
||||
|
||||
from google.adk import Agent
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.planners.plan_re_act_planner import PlanReActPlanner
|
||||
from google.adk.tools.tool_context import ToolContext
|
||||
|
||||
|
||||
def update_fc(
|
||||
data_one: str,
|
||||
data_two: Union[int, float, str],
|
||||
data_three: list[str],
|
||||
data_four: List[Union[int, float, str]],
|
||||
tool_context: ToolContext,
|
||||
) -> str:
|
||||
"""Simply ask to update these variables in the context"""
|
||||
tool_context.actions.update_state('data_one', data_one)
|
||||
tool_context.actions.update_state('data_two', data_two)
|
||||
tool_context.actions.update_state('data_three', data_three)
|
||||
tool_context.actions.update_state('data_four', data_four)
|
||||
return 'The function `update_fc` executed successfully'
|
||||
|
||||
|
||||
def echo_info(customer_id: str) -> str:
|
||||
"""Echo the context variable"""
|
||||
return customer_id
|
||||
|
||||
|
||||
def build_global_instruction(invocation_context: InvocationContext) -> str:
|
||||
return (
|
||||
'This is the global agent instruction for invocation:'
|
||||
f' {invocation_context.invocation_id}.'
|
||||
)
|
||||
|
||||
|
||||
def build_sub_agent_instruction(invocation_context: InvocationContext) -> str:
|
||||
return 'This is the plain text sub agent instruction.'
|
||||
|
||||
|
||||
context_variable_echo_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='context_variable_echo_agent',
|
||||
instruction=(
|
||||
'Use the echo_info tool to echo {customerId}, {customerInt},'
|
||||
' {customerFloat}, and {customerJson}. Ask for it if you need to.'
|
||||
),
|
||||
flow='auto',
|
||||
tools=[echo_info],
|
||||
)
|
||||
|
||||
context_variable_with_complicated_format_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='context_variable_echo_agent',
|
||||
instruction=(
|
||||
'Use the echo_info tool to echo { customerId }, {{customer_int }, { '
|
||||
" non-identifier-float}}, {artifact.fileName}, {'key1': 'value1'} and"
|
||||
" {{'key2': 'value2'}}. Ask for it if you need to."
|
||||
),
|
||||
flow='auto',
|
||||
tools=[echo_info],
|
||||
)
|
||||
|
||||
context_variable_with_nl_planner_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='context_variable_with_nl_planner_agent',
|
||||
instruction=(
|
||||
'Use the echo_info tool to echo {customerId}. Ask for it if you'
|
||||
' need to.'
|
||||
),
|
||||
flow='auto',
|
||||
planner=PlanReActPlanner(),
|
||||
tools=[echo_info],
|
||||
)
|
||||
|
||||
context_variable_with_function_instruction_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='context_variable_with_function_instruction_agent',
|
||||
instruction=build_sub_agent_instruction,
|
||||
flow='auto',
|
||||
)
|
||||
|
||||
context_variable_update_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='context_variable_update_agent',
|
||||
instruction='Call tools',
|
||||
flow='auto',
|
||||
tools=[update_fc],
|
||||
)
|
||||
|
||||
root_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='root_agent',
|
||||
description='The root agent.',
|
||||
flow='auto',
|
||||
global_instruction=build_global_instruction,
|
||||
sub_agents=[
|
||||
context_variable_with_nl_planner_agent,
|
||||
context_variable_update_agent,
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,338 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk import Agent
|
||||
|
||||
# A lightweight in-memory mock database
|
||||
ORDER_DB = {
|
||||
"1": "FINISHED",
|
||||
"2": "CANCELED",
|
||||
"3": "PENDING",
|
||||
"4": "PENDING",
|
||||
} # Order id to status mapping. Available states: 'FINISHED', 'PENDING', and 'CANCELED'
|
||||
USER_TO_ORDER_DB = {
|
||||
"user_a": ["1", "4"],
|
||||
"user_b": ["2"],
|
||||
"user_c": ["3"],
|
||||
} # User id to Order id mapping
|
||||
TICKET_DB = [{
|
||||
"ticket_id": "1",
|
||||
"user_id": "user_a",
|
||||
"issue_type": "LOGIN_ISSUE",
|
||||
"status": "OPEN",
|
||||
}] # Available states: 'OPEN', 'CLOSED', 'ESCALATED'
|
||||
USER_INFO_DB = {
|
||||
"user_a": {"name": "Alice", "email": "alice@example.com"},
|
||||
"user_b": {"name": "Bob", "email": "bob@example.com"},
|
||||
}
|
||||
|
||||
|
||||
def reset_data():
|
||||
global ORDER_DB
|
||||
global USER_TO_ORDER_DB
|
||||
global TICKET_DB
|
||||
global USER_INFO_DB
|
||||
ORDER_DB = {
|
||||
"1": "FINISHED",
|
||||
"2": "CANCELED",
|
||||
"3": "PENDING",
|
||||
"4": "PENDING",
|
||||
}
|
||||
USER_TO_ORDER_DB = {
|
||||
"user_a": ["1", "4"],
|
||||
"user_b": ["2"],
|
||||
"user_c": ["3"],
|
||||
}
|
||||
TICKET_DB = [{
|
||||
"ticket_id": "1",
|
||||
"user_id": "user_a",
|
||||
"issue_type": "LOGIN_ISSUE",
|
||||
"status": "OPEN",
|
||||
}]
|
||||
USER_INFO_DB = {
|
||||
"user_a": {"name": "Alice", "email": "alice@example.com"},
|
||||
"user_b": {"name": "Bob", "email": "bob@example.com"},
|
||||
}
|
||||
|
||||
|
||||
def get_order_status(order_id: str) -> str:
|
||||
"""Get the status of an order.
|
||||
|
||||
Args:
|
||||
order_id (str): The unique identifier of the order.
|
||||
|
||||
Returns:
|
||||
str: The status of the order (e.g., 'FINISHED', 'CANCELED', 'PENDING'),
|
||||
or 'Order not found' if the order_id does not exist.
|
||||
"""
|
||||
return ORDER_DB.get(order_id, "Order not found")
|
||||
|
||||
|
||||
def get_order_ids_for_user(user_id: str) -> list:
|
||||
"""Get the list of order IDs assigned to a specific transaction associated with a user.
|
||||
|
||||
Args:
|
||||
user_id (str): The unique identifier of the user.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of order IDs associated with the user, or an empty list
|
||||
if no orders are found.
|
||||
"""
|
||||
return USER_TO_ORDER_DB.get(user_id, [])
|
||||
|
||||
|
||||
def cancel_order(order_id: str) -> str:
|
||||
"""Cancel an order if it is in a 'PENDING' state.
|
||||
|
||||
You should call "get_order_status" to check the status first, before calling
|
||||
this tool.
|
||||
|
||||
Args:
|
||||
order_id (str): The unique identifier of the order to be canceled.
|
||||
|
||||
Returns:
|
||||
str: A message indicating whether the order was successfully canceled or
|
||||
not.
|
||||
"""
|
||||
if order_id in ORDER_DB and ORDER_DB[order_id] == "PENDING":
|
||||
ORDER_DB[order_id] = "CANCELED"
|
||||
return f"Order {order_id} has been canceled."
|
||||
return f"Order {order_id} cannot be canceled."
|
||||
|
||||
|
||||
def refund_order(order_id: str) -> str:
|
||||
"""Process a refund for an order if it is in a 'CANCELED' state.
|
||||
|
||||
You should call "get_order_status" to check if status first, before calling
|
||||
this tool.
|
||||
|
||||
Args:
|
||||
order_id (str): The unique identifier of the order to be refunded.
|
||||
|
||||
Returns:
|
||||
str: A message indicating whether the order was successfully refunded or
|
||||
not.
|
||||
"""
|
||||
if order_id in ORDER_DB and ORDER_DB[order_id] == "CANCELED":
|
||||
return f"Order {order_id} has been refunded."
|
||||
return f"Order {order_id} cannot be refunded."
|
||||
|
||||
|
||||
def create_ticket(user_id: str, issue_type: str) -> str:
|
||||
"""Create a new support ticket for a user.
|
||||
|
||||
Args:
|
||||
user_id (str): The unique identifier of the user creating the ticket.
|
||||
issue_type (str): An issue type the user is facing. Available types:
|
||||
'LOGIN_ISSUE', 'ORDER_ISSUE', 'OTHER'.
|
||||
|
||||
Returns:
|
||||
str: A message indicating that the ticket was created successfully,
|
||||
including the ticket ID.
|
||||
"""
|
||||
ticket_id = str(len(TICKET_DB) + 1)
|
||||
TICKET_DB.append({
|
||||
"ticket_id": ticket_id,
|
||||
"user_id": user_id,
|
||||
"issue_type": issue_type,
|
||||
"status": "OPEN",
|
||||
})
|
||||
return f"Ticket {ticket_id} created successfully."
|
||||
|
||||
|
||||
def get_ticket_info(ticket_id: str) -> str:
|
||||
"""Retrieve the information of a support ticket.
|
||||
|
||||
current status of a support ticket.
|
||||
|
||||
Args:
|
||||
ticket_id (str): The unique identifier of the ticket.
|
||||
|
||||
Returns:
|
||||
A dictionary contains the following fields, or 'Ticket not found' if the
|
||||
ticket_id does not exist:
|
||||
- "ticket_id": str, the current ticket id
|
||||
- "user_id": str, the associated user id
|
||||
- "issue": str, the issue type
|
||||
- "status": The current status of the ticket (e.g., 'OPEN', 'CLOSED',
|
||||
'ESCALATED')
|
||||
|
||||
Example: {"ticket_id": "1", "user_id": "user_a", "issue": "Login issue",
|
||||
"status": "OPEN"}
|
||||
"""
|
||||
for ticket in TICKET_DB:
|
||||
if ticket["ticket_id"] == ticket_id:
|
||||
return ticket
|
||||
return "Ticket not found"
|
||||
|
||||
|
||||
def get_tickets_for_user(user_id: str) -> list:
|
||||
"""Get all the ticket IDs associated with a user.
|
||||
|
||||
Args:
|
||||
user_id (str): The unique identifier of the user.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of ticket IDs associated with the user.
|
||||
If no tickets are found, returns an empty list.
|
||||
"""
|
||||
return [
|
||||
ticket["ticket_id"]
|
||||
for ticket in TICKET_DB
|
||||
if ticket["user_id"] == user_id
|
||||
]
|
||||
|
||||
|
||||
def update_ticket_status(ticket_id: str, status: str) -> str:
|
||||
"""Update the status of a support ticket.
|
||||
|
||||
Args:
|
||||
ticket_id (str): The unique identifier of the ticket.
|
||||
status (str): The new status to assign to the ticket (e.g., 'OPEN',
|
||||
'CLOSED', 'ESCALATED').
|
||||
|
||||
Returns:
|
||||
str: A message indicating whether the ticket status was successfully
|
||||
updated.
|
||||
"""
|
||||
for ticket in TICKET_DB:
|
||||
if ticket["ticket_id"] == ticket_id:
|
||||
ticket["status"] = status
|
||||
return f"Ticket {ticket_id} status updated to {status}."
|
||||
return "Ticket not found"
|
||||
|
||||
|
||||
def get_user_info(user_id: str) -> dict:
|
||||
"""Retrieve information (name, email) about a user.
|
||||
|
||||
Args:
|
||||
user_id (str): The unique identifier of the user.
|
||||
|
||||
Returns:
|
||||
dict or str: A dictionary containing user information of the following
|
||||
fields, or 'User not found' if the user_id does not exist:
|
||||
|
||||
- name: The name of the user
|
||||
- email: The email address of the user
|
||||
|
||||
For example, {"name": "Chelsea", "email": "123@example.com"}
|
||||
"""
|
||||
return USER_INFO_DB.get(user_id, "User not found")
|
||||
|
||||
|
||||
def send_email(user_id: str, email: str) -> list:
|
||||
"""Send email to user for notification.
|
||||
|
||||
Args:
|
||||
user_id (str): The unique identifier of the user.
|
||||
email (str): The email address of the user.
|
||||
|
||||
Returns:
|
||||
str: A message indicating whether the email was successfully sent.
|
||||
"""
|
||||
if user_id in USER_INFO_DB:
|
||||
return f"Email sent to {email} for user id {user_id}"
|
||||
return "Cannot find this user"
|
||||
|
||||
|
||||
# def update_user_info(user_id: str, new_info: dict[str, str]) -> str:
|
||||
def update_user_info(user_id: str, email: str, name: str) -> str:
|
||||
"""Update a user's information.
|
||||
|
||||
Args:
|
||||
user_id (str): The unique identifier of the user.
|
||||
new_info (dict): A dictionary containing the fields to be updated (e.g.,
|
||||
{'email': 'new_email@example.com'}). Available field keys: 'email' and
|
||||
'name'.
|
||||
|
||||
Returns:
|
||||
str: A message indicating whether the user's information was successfully
|
||||
updated or not.
|
||||
"""
|
||||
if user_id in USER_INFO_DB:
|
||||
# USER_INFO_DB[user_id].update(new_info)
|
||||
if email and name:
|
||||
USER_INFO_DB[user_id].update({"email": email, "name": name})
|
||||
elif email:
|
||||
USER_INFO_DB[user_id].update({"email": email})
|
||||
elif name:
|
||||
USER_INFO_DB[user_id].update({"name": name})
|
||||
else:
|
||||
raise ValueError("this should not happen.")
|
||||
return f"User {user_id} information updated."
|
||||
return "User not found"
|
||||
|
||||
|
||||
def get_user_id_from_cookie() -> str:
|
||||
"""Get user ID(username) from the cookie.
|
||||
|
||||
Only use this function when you do not know user ID(username).
|
||||
|
||||
Args: None
|
||||
|
||||
Returns:
|
||||
str: The user ID.
|
||||
"""
|
||||
return "user_a"
|
||||
|
||||
|
||||
root_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="Ecommerce_Customer_Service",
|
||||
instruction="""
|
||||
You are an intelligent customer service assistant for an e-commerce platform. Your goal is to accurately understand user queries and use the appropriate tools to fulfill requests. Follow these guidelines:
|
||||
|
||||
1. **Understand the Query**:
|
||||
- Identify actions and conditions (e.g., create a ticket only for pending orders).
|
||||
- Extract necessary details (e.g., user ID, order ID) from the query or infer them from the context.
|
||||
|
||||
2. **Plan Multi-Step Workflows**:
|
||||
- Break down complex queries into sequential steps. For example
|
||||
- typical workflow:
|
||||
- Retrieve IDs or references first (e.g., orders for a user).
|
||||
- Evaluate conditions (e.g., check order status).
|
||||
- Perform actions (e.g., create a ticket) only when conditions are met.
|
||||
- another typical workflows - order cancellation and refund:
|
||||
- Retrieve all orders for the user (`get_order_ids_for_user`).
|
||||
- Cancel pending orders (`cancel_order`).
|
||||
- Refund canceled orders (`refund_order`).
|
||||
- Notify the user (`send_email`).
|
||||
- another typical workflows - send user report:
|
||||
- Get user id.
|
||||
- Get user info(like emails)
|
||||
- Send email to user.
|
||||
|
||||
3. **Avoid Skipping Steps**:
|
||||
- Ensure each intermediate step is completed before moving to the next.
|
||||
- Do not create tickets or take other actions without verifying the conditions specified in the query.
|
||||
|
||||
4. **Provide Clear Responses**:
|
||||
- Confirm the actions performed, including details like ticket ID or pending orders.
|
||||
- Ensure the response aligns with the steps taken and query intent.
|
||||
""",
|
||||
tools=[
|
||||
get_order_status,
|
||||
cancel_order,
|
||||
get_order_ids_for_user,
|
||||
refund_order,
|
||||
create_ticket,
|
||||
update_ticket_status,
|
||||
get_tickets_for_user,
|
||||
get_ticket_info,
|
||||
get_user_info,
|
||||
send_email,
|
||||
update_user_info,
|
||||
get_user_id_from_cookie,
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,229 @@
|
||||
{
|
||||
"eval_set_id": "a1157c01-851f-48a8-b956-83cf7f463510",
|
||||
"name": "a1157c01-851f-48a8-b956-83cf7f463510",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/ecommerce_customer_service_agent/order_query.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "38d54523-d789-4873-8cc0-d38826c7feb4",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Send an email to user user_a whose email address is alice@example.com"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Email sent to alice@example.com for user id user_a."
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"email": "alice@example.com",
|
||||
"user_id": "user_a"
|
||||
},
|
||||
"name": "send_email"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747341706.6240807
|
||||
},
|
||||
{
|
||||
"invocation_id": "916393ab-0bce-4cb0-98de-6573d4e8e25c",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Can you tell me the status of my order with ID 1?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Your order with ID 1 is FINISHED."
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"order_id": "1"
|
||||
},
|
||||
"name": "get_order_status"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747341706.6241167
|
||||
},
|
||||
{
|
||||
"invocation_id": "511b23d9-56f9-423b-9c31-7626f3411c32",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Cancel all pending order for the user with user id user_a"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "I have checked your orders and order 4 was in pending status, so I have cancelled it. Order 1 was already finished and couldn't be cancelled.\n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"user_id": "user_a"
|
||||
},
|
||||
"name": "get_order_ids_for_user"
|
||||
},
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"order_id": "1"
|
||||
},
|
||||
"name": "get_order_status"
|
||||
},
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"order_id": "4"
|
||||
},
|
||||
"name": "get_order_status"
|
||||
},
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"order_id": "4"
|
||||
},
|
||||
"name": "cancel_order"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747341706.6241703
|
||||
},
|
||||
{
|
||||
"invocation_id": "dcdf4b6d-96dd-4602-8c14-0563c6f6b5d0",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "What orders have I placed under the username user_b?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "User user_b has placed one order with order ID 2.\n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"user_id": "user_b"
|
||||
},
|
||||
"name": "get_order_ids_for_user"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747341706.624196
|
||||
}
|
||||
],
|
||||
"session_input": null,
|
||||
"creation_timestamp": 1747341706.6242023
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747341706.6242158
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"criteria": {
|
||||
"tool_trajectory_avg_score": 0.7,
|
||||
"response_match_score": 0.5
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,182 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk import Agent
|
||||
from google.genai import types
|
||||
|
||||
research_plan_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="research_plan_agent",
|
||||
description="I can help generate research plan.",
|
||||
instruction="""\
|
||||
Your task is to create a research plan according to the user's query.
|
||||
|
||||
# Here are the instructions for creating the research plan:
|
||||
|
||||
+ Focus on finding specific things, e.g. products, data, etc.
|
||||
+ Have the personality of a work colleague that is very helpful and explains things very nicely.
|
||||
+ Don't mention your name unless you are asked.
|
||||
+ Think about the most common things that you would need to research.
|
||||
+ Think about possible answers when creating the plan.
|
||||
+ Your task is to create the sections that should be researched. You will output high level headers, preceded by ##
|
||||
+ Underneath each header, write a short sentence on what we want to find there.
|
||||
+ The headers will follow the logical analysis pattern, as well as logical exploration pattern.
|
||||
+ The headers should be a statement, not be in the form of questions.
|
||||
+ The header will not include roman numerals or anything of the sort, e.g. ":", etc
|
||||
+ Do not include things that you cannot possibly know about from using Google Search: e.g. sales forecasting, competitors, profitability analysis, etc.
|
||||
+ Do not have an executive summary
|
||||
+ In each section describe specifically what will be researched.
|
||||
+ Never use "we will", but rather "I will".
|
||||
+ Don't ask for clarifications from the user.
|
||||
+ Do not ask the user for clarifications or if they have any other questions.
|
||||
+ All headers should be bolded.
|
||||
+ If you have steps in the plan that depend on other information, make sure they are 2 different sections in the plan.
|
||||
+ At the end mention that you will start researching.
|
||||
|
||||
# Instruction on replying format
|
||||
|
||||
+ Start with your name as "[research_plan_agent]: ".
|
||||
+ Output the content you want to say.
|
||||
|
||||
Output summary:
|
||||
""",
|
||||
flow="single",
|
||||
sub_agents=[],
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
question_generation_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="question_generation_agent",
|
||||
description="I can help generate questions related to user's question.",
|
||||
instruction="""\
|
||||
Generate questions related to the research plan generated by research_plan_agent.
|
||||
|
||||
# Instruction on replying format
|
||||
|
||||
Your reply should be a numbered list.
|
||||
|
||||
For each question, reply in the following format: "[question_generation_agent]: [generated questions]"
|
||||
|
||||
Here is an example of the generated question list:
|
||||
|
||||
1. [question_generation_agent]: which state is San Jose in?
|
||||
2. [question_generation_agent]: how google website is designed?
|
||||
""",
|
||||
flow="single",
|
||||
sub_agents=[],
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
),
|
||||
)
|
||||
|
||||
information_retrieval_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="information_retrieval_agent",
|
||||
description=(
|
||||
"I can help retrieve information related to question_generation_agent's"
|
||||
" question."
|
||||
),
|
||||
instruction="""\
|
||||
Inspect all the questions after "[question_generation_agent]: " and answer them.
|
||||
|
||||
# Instruction on replying format
|
||||
|
||||
Always start with "[information_retrieval_agent]: "
|
||||
|
||||
For the answer of one question:
|
||||
|
||||
- Start with a title with one line summary of the reply.
|
||||
- The title line should be bolded and starts with No.x of the corresponding question.
|
||||
- Have a paragraph of detailed explain.
|
||||
|
||||
# Instruction on exiting the loop
|
||||
|
||||
- If you see there are less than 20 questions by "question_generation_agent", do not say "[exit]".
|
||||
- If you see there are already great or equal to 20 questions asked by "question_generation_agent", say "[exit]" at last to exit the loop.
|
||||
""",
|
||||
flow="single",
|
||||
sub_agents=[],
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
),
|
||||
)
|
||||
|
||||
question_sources_generation_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="question_sources_generation_agent",
|
||||
description=(
|
||||
"I can help generate questions and retrieve related information."
|
||||
),
|
||||
instruction="Generate questions and retrieve information.",
|
||||
flow="loop",
|
||||
sub_agents=[
|
||||
question_generation_agent,
|
||||
information_retrieval_agent,
|
||||
],
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
),
|
||||
)
|
||||
|
||||
summary_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="summary_agent",
|
||||
description="I can help summarize information of previous content.",
|
||||
instruction="""\
|
||||
Summarize information in all historical messages that were replied by "question_generation_agent" and "information_retrieval_agent".
|
||||
|
||||
# Instruction on replying format
|
||||
|
||||
- The output should be like an essay that has a title, an abstract, multiple paragraphs for each topic and a conclusion.
|
||||
- Each paragraph should maps to one or more question in historical content.
|
||||
""",
|
||||
flow="single",
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.8,
|
||||
),
|
||||
)
|
||||
|
||||
research_assistant = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="research_assistant",
|
||||
description="I can help with research question.",
|
||||
instruction="Help customers with their need.",
|
||||
flow="sequential",
|
||||
sub_agents=[
|
||||
research_plan_agent,
|
||||
question_sources_generation_agent,
|
||||
summary_agent,
|
||||
],
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
),
|
||||
)
|
||||
|
||||
spark_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="spark_assistant",
|
||||
description="I can help with non-research question.",
|
||||
instruction="Help customers with their need.",
|
||||
flow="auto",
|
||||
sub_agents=[research_assistant],
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
),
|
||||
)
|
||||
|
||||
root_agent = spark_agent
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,95 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
# Hello world agent from agent 1.0 - https://colab.sandbox.google.com/drive/1Zq-nqmgK0nCERCv8jKIaoeTTgbNn6oSo?resourcekey=0-GYaz9pFT4wY8CI8Cvjy5GA#scrollTo=u3X3XwDOaCv9
|
||||
import random
|
||||
|
||||
from google.adk import Agent
|
||||
from google.genai import types
|
||||
|
||||
|
||||
def roll_die(sides: int) -> int:
|
||||
"""Roll a die and return the rolled result.
|
||||
|
||||
Args:
|
||||
sides: The integer number of sides the die has.
|
||||
|
||||
Returns:
|
||||
An integer of the result of rolling the die.
|
||||
"""
|
||||
return random.randint(1, sides)
|
||||
|
||||
|
||||
def check_prime(nums: list[int]) -> list[str]:
|
||||
"""Check if a given list of numbers are prime.
|
||||
|
||||
Args:
|
||||
nums: The list of numbers to check.
|
||||
|
||||
Returns:
|
||||
A str indicating which number is prime.
|
||||
"""
|
||||
primes = set()
|
||||
for number in nums:
|
||||
number = int(number)
|
||||
if number <= 1:
|
||||
continue
|
||||
is_prime = True
|
||||
for i in range(2, int(number**0.5) + 1):
|
||||
if number % i == 0:
|
||||
is_prime = False
|
||||
break
|
||||
if is_prime:
|
||||
primes.add(number)
|
||||
return (
|
||||
'No prime numbers found.'
|
||||
if not primes
|
||||
else f"{', '.join(str(num) for num in primes)} are prime numbers."
|
||||
)
|
||||
|
||||
|
||||
root_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='data_processing_agent',
|
||||
instruction="""
|
||||
You roll dice and answer questions about the outcome of the dice rolls.
|
||||
You can roll dice of different sizes.
|
||||
You can use multiple tools in parallel by calling functions in parallel(in one request and in one round).
|
||||
The only things you do are roll dice for the user and discuss the outcomes.
|
||||
It is ok to discuss previous dice roles, and comment on the dice rolls.
|
||||
When you are asked to roll a die, you must call the roll_die tool with the number of sides. Be sure to pass in an integer. Do not pass in a string.
|
||||
You should never roll a die on your own.
|
||||
When checking prime numbers, call the check_prime tool with a list of integers. Be sure to pass in a list of integers. You should never pass in a string.
|
||||
You should not check prime numbers before calling the tool.
|
||||
When you are asked to roll a die and check prime numbers, you should always make the following two function calls:
|
||||
1. You should first call the roll_die tool to get a roll. Wait for the function response before calling the check_prime tool.
|
||||
2. After you get the function response from roll_die tool, you should call the check_prime tool with the roll_die result.
|
||||
2.1 If user asks you to check primes based on previous rolls, make sure you include the previous rolls in the list.
|
||||
3. When you respond, you must include the roll_die result from step 1.
|
||||
You should always perform the previous 3 steps when asking for a roll and checking prime numbers.
|
||||
You should not rely on the previous history on prime results.
|
||||
""",
|
||||
tools=[
|
||||
roll_die,
|
||||
check_prime,
|
||||
],
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
safety_settings=[
|
||||
types.SafetySetting( # avoid false alarm about rolling dice.
|
||||
category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
|
||||
threshold=types.HarmBlockThreshold.OFF,
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,143 @@
|
||||
{
|
||||
"eval_set_id": "56540925-a5ff-49fe-a4e1-589fe78066f2",
|
||||
"name": "56540925-a5ff-49fe-a4e1-589fe78066f2",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/hello_world_agent/roll_die.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "b01f67f0-9f23-44d6-bbe4-36ea235cb9fb",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Hi who are you?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "I am a data processing agent. I can roll dice and check if the results are prime numbers. What would you like me to do? \n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747341775.8937013
|
||||
},
|
||||
{
|
||||
"invocation_id": "13be0093-ac29-4828-98c6-5bbd570c010c",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "What can you do?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "I can roll dice for you of different sizes, and I can check if the results are prime numbers. I can also remember previous rolls if you'd like to check those for primes as well. What would you like me to do? \n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747341775.8937378
|
||||
},
|
||||
{
|
||||
"invocation_id": "7deda353-c936-4c21-b242-9fa75e45b6a7",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Can you roll a die with 6 sides"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": null
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"sides": 6
|
||||
},
|
||||
"name": "roll_die"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747341775.8937788
|
||||
}
|
||||
],
|
||||
"session_input": null,
|
||||
"creation_timestamp": 1747341775.8937826
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747341775.8937957
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"criteria": {
|
||||
"tool_trajectory_avg_score": 1.0,
|
||||
"response_match_score": 0.5
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,104 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
# Hello world agent from agent 1.0 revised to be defined with get_agent_async
|
||||
# instead of root_agent - https://colab.sandbox.google.com/drive/1Zq-nqmgK0nCERCv8jKIaoeTTgbNn6oSo?resourcekey=0-GYaz9pFT4wY8CI8Cvjy5GA#scrollTo=u3X3XwDOaCv9
|
||||
import contextlib
|
||||
import random
|
||||
from typing import Optional
|
||||
|
||||
from google.adk import Agent
|
||||
from google.adk.agents import llm_agent
|
||||
from google.genai import types
|
||||
|
||||
|
||||
def roll_die(sides: int) -> int:
|
||||
"""Roll a die and return the rolled result.
|
||||
|
||||
Args:
|
||||
sides: The integer number of sides the die has.
|
||||
|
||||
Returns:
|
||||
An integer of the result of rolling the die.
|
||||
"""
|
||||
return random.randint(1, sides)
|
||||
|
||||
|
||||
def check_prime(nums: list[int]) -> list[str]:
|
||||
"""Check if a given list of numbers are prime.
|
||||
|
||||
Args:
|
||||
nums: The list of numbers to check.
|
||||
|
||||
Returns:
|
||||
A str indicating which number is prime.
|
||||
"""
|
||||
primes = set()
|
||||
for number in nums:
|
||||
number = int(number)
|
||||
if number <= 1:
|
||||
continue
|
||||
is_prime = True
|
||||
for i in range(2, int(number**0.5) + 1):
|
||||
if number % i == 0:
|
||||
is_prime = False
|
||||
break
|
||||
if is_prime:
|
||||
primes.add(number)
|
||||
return (
|
||||
'No prime numbers found.'
|
||||
if not primes
|
||||
else f"{', '.join(str(num) for num in primes)} are prime numbers."
|
||||
)
|
||||
|
||||
|
||||
async def get_agent_async() -> (
|
||||
tuple[llm_agent.LlmAgent, Optional[contextlib.AsyncExitStack]]
|
||||
):
|
||||
"""Returns the root agent."""
|
||||
root_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='data_processing_agent',
|
||||
instruction="""
|
||||
You roll dice and answer questions about the outcome of the dice rolls.
|
||||
You can roll dice of different sizes.
|
||||
You can use multiple tools in parallel by calling functions in parallel(in one request and in one round).
|
||||
The only things you do are roll dice for the user and discuss the outcomes.
|
||||
It is ok to discuss previous dice roles, and comment on the dice rolls.
|
||||
When you are asked to roll a die, you must call the roll_die tool with the number of sides. Be sure to pass in an integer. Do not pass in a string.
|
||||
You should never roll a die on your own.
|
||||
When checking prime numbers, call the check_prime tool with a list of integers. Be sure to pass in a list of integers. You should never pass in a string.
|
||||
You should not check prime numbers before calling the tool.
|
||||
When you are asked to roll a die and check prime numbers, you should always make the following two function calls:
|
||||
1. You should first call the roll_die tool to get a roll. Wait for the function response before calling the check_prime tool.
|
||||
2. After you get the function response from roll_die tool, you should call the check_prime tool with the roll_die result.
|
||||
2.1 If user asks you to check primes based on previous rolls, make sure you include the previous rolls in the list.
|
||||
3. When you respond, you must include the roll_die result from step 1.
|
||||
You should always perform the previous 3 steps when asking for a roll and checking prime numbers.
|
||||
You should not rely on the previous history on prime results.
|
||||
""",
|
||||
tools=[
|
||||
roll_die,
|
||||
check_prime,
|
||||
],
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
safety_settings=[
|
||||
types.SafetySetting( # avoid false alarm about rolling dice.
|
||||
category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
|
||||
threshold=types.HarmBlockThreshold.OFF,
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
return root_agent, None
|
||||
@@ -0,0 +1,55 @@
|
||||
{
|
||||
"eval_set_id": "56540925-a5ff-49fe-a4e1-589fe78066f2",
|
||||
"name": "56540925-a5ff-49fe-a4e1-589fe78066f2",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/hello_world_agent_async/roll_die.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "b01f67f0-9f23-44d6-bbe4-36ea235cb9fb",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Hi who are you?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "I am a data processing agent. I can roll dice and check if the results are prime numbers. What would you like me to do? \n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747341775.8937013
|
||||
}
|
||||
],
|
||||
"session_input": null,
|
||||
"creation_timestamp": 1747341775.8937826
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747341775.8937957
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"criteria": {
|
||||
"tool_trajectory_avg_score": 1.0,
|
||||
"response_match_score": 0.5
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,304 @@
|
||||
# Copyright 2026 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 os
|
||||
import sys
|
||||
|
||||
from google.adk import Agent
|
||||
|
||||
DEVICE_DB = {
|
||||
"device_1": {"status": "ON", "location": "Living Room"},
|
||||
"device_2": {"status": "OFF", "location": "Bedroom"},
|
||||
"device_3": {"status": "OFF", "location": "Kitchen"},
|
||||
}
|
||||
|
||||
TEMPERATURE_DB = {
|
||||
"Living Room": 22,
|
||||
"Bedroom": 20,
|
||||
"Kitchen": 24,
|
||||
}
|
||||
|
||||
SCHEDULE_DB = {
|
||||
"device_1": {"time": "18:00", "status": "ON"},
|
||||
"device_2": {"time": "22:00", "status": "OFF"},
|
||||
}
|
||||
|
||||
USER_PREFERENCES_DB = {
|
||||
"user_x": {"preferred_temp": 21, "location": "Bedroom"},
|
||||
"user_x": {"preferred_temp": 21, "location": "Living Room"},
|
||||
"user_y": {"preferred_temp": 23, "location": "Living Room"},
|
||||
}
|
||||
|
||||
|
||||
def reset_data():
|
||||
global DEVICE_DB
|
||||
global TEMPERATURE_DB
|
||||
global SCHEDULE_DB
|
||||
global USER_PREFERENCES_DB
|
||||
DEVICE_DB = {
|
||||
"device_1": {"status": "ON", "location": "Living Room"},
|
||||
"device_2": {"status": "OFF", "location": "Bedroom"},
|
||||
"device_3": {"status": "OFF", "location": "Kitchen"},
|
||||
}
|
||||
|
||||
TEMPERATURE_DB = {
|
||||
"Living Room": 22,
|
||||
"Bedroom": 20,
|
||||
"Kitchen": 24,
|
||||
}
|
||||
|
||||
SCHEDULE_DB = {
|
||||
"device_1": {"time": "18:00", "status": "ON"},
|
||||
"device_2": {"time": "22:00", "status": "OFF"},
|
||||
}
|
||||
|
||||
USER_PREFERENCES_DB = {
|
||||
"user_x": {"preferred_temp": 21, "location": "Bedroom"},
|
||||
"user_x": {"preferred_temp": 21, "location": "Living Room"},
|
||||
"user_y": {"preferred_temp": 23, "location": "Living Room"},
|
||||
}
|
||||
|
||||
|
||||
def get_device_info(device_id: str) -> dict:
|
||||
"""Get the current status and location of a AC device.
|
||||
|
||||
Args:
|
||||
device_id (str): The unique identifier of the device.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the following fields, or 'Device not found'
|
||||
if the device_id does not exist:
|
||||
- status: The current status of the device (e.g., 'ON', 'OFF')
|
||||
- location: The location where the device is installed (e.g., 'Living
|
||||
Room', 'Bedroom', ''Kitchen')
|
||||
"""
|
||||
return DEVICE_DB.get(device_id, "Device not found")
|
||||
|
||||
|
||||
# def set_device_info(device_id: str, updates: dict) -> str:
|
||||
# """Update the information of a AC device, specifically its status and/or location.
|
||||
|
||||
# Args:
|
||||
# device_id (str): Required. The unique identifier of the device.
|
||||
# updates (dict): Required. A dictionary containing the fields to be
|
||||
# updated. Supported keys: - "status" (str): The new status to set for the
|
||||
# device. Accepted values: 'ON', 'OFF'. **Only these values are allowed.**
|
||||
# - "location" (str): The new location to set for the device. Accepted
|
||||
# values: 'Living Room', 'Bedroom', 'Kitchen'. **Only these values are
|
||||
# allowed.**
|
||||
|
||||
|
||||
# Returns:
|
||||
# str: A message indicating whether the device information was successfully
|
||||
# updated.
|
||||
# """
|
||||
# if device_id in DEVICE_DB:
|
||||
# if "status" in updates:
|
||||
# DEVICE_DB[device_id]["status"] = updates["status"]
|
||||
# if "location" in updates:
|
||||
# DEVICE_DB[device_id]["location"] = updates["location"]
|
||||
# return f"Device {device_id} information updated: {updates}."
|
||||
# return "Device not found"
|
||||
def set_device_info(
|
||||
device_id: str, status: str = "", location: str = ""
|
||||
) -> str:
|
||||
"""Update the information of a AC device, specifically its status and/or location.
|
||||
|
||||
Args:
|
||||
device_id (str): Required. The unique identifier of the device.
|
||||
status (str): The new status to set for the
|
||||
device. Accepted values: 'ON', 'OFF'. **Only these values are allowed.**
|
||||
location (str): The new location to set for the device. Accepted
|
||||
values: 'Living Room', 'Bedroom', 'Kitchen'. **Only these values are
|
||||
allowed.**
|
||||
|
||||
Returns:
|
||||
str: A message indicating whether the device information was successfully
|
||||
updated.
|
||||
"""
|
||||
if device_id in DEVICE_DB:
|
||||
if status:
|
||||
DEVICE_DB[device_id]["status"] = status
|
||||
return f"Device {device_id} information updated: status -> {status}."
|
||||
if location:
|
||||
DEVICE_DB[device_id]["location"] = location
|
||||
return f"Device {device_id} information updated: location -> {location}."
|
||||
return "Device not found"
|
||||
|
||||
|
||||
def get_temperature(location: str) -> int:
|
||||
"""Get the current temperature in Celsius of a location (e.g., 'Living Room', 'Bedroom', 'Kitchen').
|
||||
|
||||
Args:
|
||||
location (str): The location for which to retrieve the temperature (e.g.,
|
||||
'Living Room', 'Bedroom', 'Kitchen').
|
||||
|
||||
Returns:
|
||||
int: The current temperature in Celsius in the specified location, or
|
||||
'Location not found' if the location does not exist.
|
||||
"""
|
||||
return TEMPERATURE_DB.get(location, "Location not found")
|
||||
|
||||
|
||||
def set_temperature(location: str, temperature: int) -> str:
|
||||
"""Set the desired temperature in Celsius for a location.
|
||||
|
||||
Acceptable range of temperature: 18-30 Celsius. If it's out of the range, do
|
||||
not call this tool.
|
||||
|
||||
Args:
|
||||
location (str): The location where the temperature should be set.
|
||||
temperature (int): The desired temperature as integer to set in Celsius.
|
||||
Acceptable range: 18-30 Celsius.
|
||||
|
||||
Returns:
|
||||
str: A message indicating whether the temperature was successfully set.
|
||||
"""
|
||||
if location in TEMPERATURE_DB:
|
||||
TEMPERATURE_DB[location] = temperature
|
||||
return f"Temperature in {location} set to {temperature}°C."
|
||||
return "Location not found"
|
||||
|
||||
|
||||
def get_user_preferences(user_id: str) -> dict:
|
||||
"""Get the temperature preferences and preferred location of a user_id.
|
||||
|
||||
user_id must be provided.
|
||||
|
||||
Args:
|
||||
user_id (str): The unique identifier of the user.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the following fields, or 'User not found' if
|
||||
the user_id does not exist:
|
||||
- preferred_temp: The user's preferred temperature.
|
||||
- location: The location where the user prefers to be.
|
||||
"""
|
||||
return USER_PREFERENCES_DB.get(user_id, "User not found")
|
||||
|
||||
|
||||
def set_device_schedule(device_id: str, time: str, status: str) -> str:
|
||||
"""Schedule a device to change its status at a specific time.
|
||||
|
||||
Args:
|
||||
device_id (str): The unique identifier of the device.
|
||||
time (str): The time at which the device should change its status (format:
|
||||
'HH:MM').
|
||||
status (str): The status to set for the device at the specified time
|
||||
(e.g., 'ON', 'OFF').
|
||||
|
||||
Returns:
|
||||
str: A message indicating whether the schedule was successfully set.
|
||||
"""
|
||||
if device_id in DEVICE_DB:
|
||||
SCHEDULE_DB[device_id] = {"time": time, "status": status}
|
||||
return f"Device {device_id} scheduled to turn {status} at {time}."
|
||||
return "Device not found"
|
||||
|
||||
|
||||
def get_device_schedule(device_id: str) -> dict:
|
||||
"""Retrieve the schedule of a device.
|
||||
|
||||
Args:
|
||||
device_id (str): The unique identifier of the device.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the following fields, or 'Schedule not
|
||||
found' if the device_id does not exist:
|
||||
- time: The scheduled time for the device to change its status (format:
|
||||
'HH:MM').
|
||||
- status: The status that will be set at the scheduled time (e.g., 'ON',
|
||||
'OFF').
|
||||
"""
|
||||
return SCHEDULE_DB.get(device_id, "Schedule not found")
|
||||
|
||||
|
||||
def celsius_to_fahrenheit(celsius: int) -> float:
|
||||
"""Convert Celsius to Fahrenheit.
|
||||
|
||||
You must call this to do the conversion of temperature, so you can get the
|
||||
precise number in required format.
|
||||
|
||||
Args:
|
||||
celsius (int): Temperature in Celsius.
|
||||
|
||||
Returns:
|
||||
float: Temperature in Fahrenheit.
|
||||
"""
|
||||
return (celsius * 9 / 5) + 32
|
||||
|
||||
|
||||
def fahrenheit_to_celsius(fahrenheit: float) -> int:
|
||||
"""Convert Fahrenheit to Celsius.
|
||||
|
||||
You must call this to do the conversion of temperature, so you can get the
|
||||
precise number in required format.
|
||||
|
||||
Args:
|
||||
fahrenheit (float): Temperature in Fahrenheit.
|
||||
|
||||
Returns:
|
||||
int: Temperature in Celsius.
|
||||
"""
|
||||
return int((fahrenheit - 32) * 5 / 9)
|
||||
|
||||
|
||||
def list_devices(status: str = "", location: str = "") -> list:
|
||||
"""Retrieve a list of AC devices, filtered by status and/or location when provided.
|
||||
|
||||
For cost efficiency, always apply as many filters (status and location) as
|
||||
available in the input arguments.
|
||||
|
||||
Args:
|
||||
status (str, optional): The status to filter devices by (e.g., 'ON',
|
||||
'OFF'). Defaults to None.
|
||||
location (str, optional): The location to filter devices by (e.g., 'Living
|
||||
Room', 'Bedroom', ''Kitchen'). Defaults to None.
|
||||
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing the device ID, status, and
|
||||
location, or an empty list if no devices match the criteria.
|
||||
"""
|
||||
devices = []
|
||||
for device_id, info in DEVICE_DB.items():
|
||||
if ((not status) or info["status"] == status) and (
|
||||
(not location) or info["location"] == location
|
||||
):
|
||||
devices.append({
|
||||
"device_id": device_id,
|
||||
"status": info["status"],
|
||||
"location": info["location"],
|
||||
})
|
||||
return devices if devices else "No devices found matching the criteria."
|
||||
|
||||
|
||||
root_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="Home_automation_agent",
|
||||
instruction="""
|
||||
You are Home Automation Agent. You are responsible for controlling the devices in the home.
|
||||
""",
|
||||
tools=[
|
||||
get_device_info,
|
||||
set_device_info,
|
||||
get_temperature,
|
||||
set_temperature,
|
||||
get_user_preferences,
|
||||
set_device_schedule,
|
||||
get_device_schedule,
|
||||
celsius_to_fahrenheit,
|
||||
fahrenheit_to_celsius,
|
||||
list_devices,
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,65 @@
|
||||
{
|
||||
"eval_set_id": "b305bd06-38c5-4796-b9c7-d9c7454338b9",
|
||||
"name": "b305bd06-38c5-4796-b9c7-d9c7454338b9",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/home_automation_agent/simple_test.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "b7982664-0ab6-47cc-ab13-326656afdf75",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Turn off device_2 in the Bedroom."
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "I have set the device_2 status to off."
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"location": "Bedroom",
|
||||
"device_id": "device_2",
|
||||
"status": "OFF"
|
||||
},
|
||||
"name": "set_device_info"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747337309.2360144
|
||||
}
|
||||
],
|
||||
"session_input": null,
|
||||
"creation_timestamp": 1747337309.2360282
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747337309.2360387
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
[{
|
||||
"query": "Turn off device_3 in the Bedroom.",
|
||||
"expected_tool_use": [{"tool_name": "set_device_info", "tool_input": {"location": "Bedroom", "device_id": "device_3", "status": "OFF"}}],
|
||||
"reference": "I have set the device_3 status to off."
|
||||
}]
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"criteria": {
|
||||
"tool_trajectory_avg_score": 1.0,
|
||||
"response_match_score": 0.3
|
||||
}
|
||||
}
|
||||
+113
@@ -0,0 +1,113 @@
|
||||
{
|
||||
"eval_set_id": "1be50511-ff75-4d68-b2d7-2165cbdc1044",
|
||||
"name": "1be50511-ff75-4d68-b2d7-2165cbdc1044",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/home_automation_agent/test_files/dependent_tool_calls.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "cbece1c0-3811-45c0-96fc-9a4279075483",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Turn off device_2 in the Bedroom."
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "I have set the device 2 status to off."
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"location": "Bedroom",
|
||||
"status": "OFF",
|
||||
"device_id": "device_2"
|
||||
},
|
||||
"name": "set_device_info"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747340826.1082227
|
||||
},
|
||||
{
|
||||
"invocation_id": "cc85cdae-4258-4b94-8fe7-a985b8356190",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "What's the status of device_2 in the Bedroom?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Status of device_2 is off."
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"device_id": "device_2"
|
||||
},
|
||||
"name": "get_device_info"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747340826.1082554
|
||||
}
|
||||
],
|
||||
"session_input": null,
|
||||
"creation_timestamp": 1747340826.108262
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747340826.108275
|
||||
}
|
||||
+105
@@ -0,0 +1,105 @@
|
||||
{
|
||||
"eval_set_id": "94553685-5f19-492b-bc44-f3bc775955e9",
|
||||
"name": "94553685-5f19-492b-bc44-f3bc775955e9",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/home_automation_agent/test_files/memorizing_past_events/eval_data.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "a958b622-21d3-4a6c-9c15-1274bbb8a6b6",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Turn off device_2 in the Bedroom."
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "OK. I've turned off device_2 in the Bedroom. Anything else?\n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"location": "Bedroom",
|
||||
"device_id": "device_2",
|
||||
"status": "OFF"
|
||||
},
|
||||
"name": "set_device_info"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747340865.7043095
|
||||
},
|
||||
{
|
||||
"invocation_id": "1c07123d-4bed-4eb0-9e55-c7f80c70dadf",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "What's the command I just issued?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "You asked me to turn off device_2 in the Bedroom.\n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747340865.7043421
|
||||
}
|
||||
],
|
||||
"session_input": null,
|
||||
"creation_timestamp": 1747340865.7043483
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747340865.704361
|
||||
}
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"criteria": {
|
||||
"tool_trajectory_avg_score": 1.0,
|
||||
"response_match_score": 0.5
|
||||
}
|
||||
}
|
||||
+115
@@ -0,0 +1,115 @@
|
||||
{
|
||||
"eval_set_id": "4412cca6-dfcd-43ab-bbc5-9155380c7137",
|
||||
"name": "4412cca6-dfcd-43ab-bbc5-9155380c7137",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/home_automation_agent/test_files/simple_multi_turn_conversation.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "9f51a1ac-56a4-4b4a-9878-36ff1ae312ce",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Turn off device_2 in the Bedroom."
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "I have set the device 2 status to off."
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"location": "Bedroom",
|
||||
"device_id": "device_2",
|
||||
"status": "OFF"
|
||||
},
|
||||
"name": "set_device_info"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747340791.7353904
|
||||
},
|
||||
{
|
||||
"invocation_id": "c82d54d0-5fa8-4f79-a6dc-692090f0d42b",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Turn on device_2 in the Bedroom."
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "I have set the device 2 status to on."
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"location": "Bedroom",
|
||||
"status": "ON",
|
||||
"device_id": "device_2"
|
||||
},
|
||||
"name": "set_device_info"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747340791.7354295
|
||||
}
|
||||
],
|
||||
"session_input": null,
|
||||
"creation_timestamp": 1747340791.7354348
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747340791.735446
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
{
|
||||
"eval_set_id": "9100bfc9-cc28-4ab9-b920-2dc72e138997",
|
||||
"name": "9100bfc9-cc28-4ab9-b920-2dc72e138997",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/home_automation_agent/test_files/simple_test.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "9f5e8d91-8e51-41d6-addf-196a828168c5",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Turn off device_2 in the Bedroom."
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "OK. I've turned off device_2 in the Bedroom. Anything else?\n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"location": "Bedroom",
|
||||
"device_id": "device_2",
|
||||
"status": "OFF"
|
||||
},
|
||||
"name": "set_device_info"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747340849.0429707
|
||||
},
|
||||
{
|
||||
"invocation_id": "767b2451-5f7b-4c73-aeaf-a82c71e15788",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "What's the command I just issued?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "You asked me to turn off device_2 in the Bedroom.\n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747340849.0429986
|
||||
}
|
||||
],
|
||||
"session_input": null,
|
||||
"creation_timestamp": 1747340849.0430045
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747340849.0430162
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
{
|
||||
"eval_set_id": "e141f90b-9e7e-4f06-94d7-bbe7e8080ead",
|
||||
"name": "e141f90b-9e7e-4f06-94d7-bbe7e8080ead",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/home_automation_agent/test_files/simple_test2.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "c35582f7-838a-460f-b783-039e278165e0",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Turn off device_3 in the Bedroom."
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "I have set the device_3 status to off."
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"location": "Bedroom",
|
||||
"device_id": "device_3",
|
||||
"status": "OFF"
|
||||
},
|
||||
"name": "set_device_info"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747340814.8645504
|
||||
}
|
||||
],
|
||||
"session_input": null,
|
||||
"creation_timestamp": 1747340814.86456
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747340814.864572
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"criteria": {
|
||||
"tool_trajectory_avg_score": 1.0
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,218 @@
|
||||
# Copyright 2026 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 os
|
||||
from typing import Any
|
||||
|
||||
from crewai_tools import DirectoryReadTool
|
||||
from google.adk import Agent
|
||||
from google.adk.tools.agent_tool import AgentTool
|
||||
from google.adk.tools.crewai_tool import CrewaiTool
|
||||
from google.adk.tools.langchain_tool import LangchainTool
|
||||
from google.adk.tools.retrieval.files_retrieval import FilesRetrieval
|
||||
from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
|
||||
from langchain_community.tools import ShellTool
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class TestCase(BaseModel):
|
||||
case: str
|
||||
|
||||
|
||||
class Test(BaseModel):
|
||||
test_title: list[str]
|
||||
|
||||
|
||||
def simple_function(param: str) -> str:
|
||||
if isinstance(param, str):
|
||||
return "Called simple function successfully"
|
||||
return "Called simple function with wrong param type"
|
||||
|
||||
|
||||
def no_param_function() -> str:
|
||||
return "Called no param function successfully"
|
||||
|
||||
|
||||
def no_output_function(param: str):
|
||||
return
|
||||
|
||||
|
||||
def multiple_param_types_function(
|
||||
param1: str, param2: int, param3: float, param4: bool
|
||||
) -> str:
|
||||
if (
|
||||
isinstance(param1, str)
|
||||
and isinstance(param2, int)
|
||||
and isinstance(param3, float)
|
||||
and isinstance(param4, bool)
|
||||
):
|
||||
return "Called multiple param types function successfully"
|
||||
return "Called multiple param types function with wrong param types"
|
||||
|
||||
|
||||
def throw_error_function(param: str) -> str:
|
||||
raise ValueError("Error thrown by throw_error_function")
|
||||
|
||||
|
||||
def list_str_param_function(param: list[str]) -> str:
|
||||
if isinstance(param, list) and all(isinstance(item, str) for item in param):
|
||||
return "Called list str param function successfully"
|
||||
return "Called list str param function with wrong param type"
|
||||
|
||||
|
||||
def return_list_str_function(param: str) -> list[str]:
|
||||
return ["Called return list str function successfully"]
|
||||
|
||||
|
||||
def complex_function_list_dict(
|
||||
param1: dict[str, Any], param2: list[dict[str, Any]]
|
||||
) -> list[Test]:
|
||||
if (
|
||||
isinstance(param1, dict)
|
||||
and isinstance(param2, list)
|
||||
and all(isinstance(item, dict) for item in param2)
|
||||
):
|
||||
return [
|
||||
Test(test_title=["function test 1", "function test 2"]),
|
||||
Test(test_title=["retrieval test"]),
|
||||
]
|
||||
raise ValueError("Wrong param")
|
||||
|
||||
|
||||
def repetitive_call_1(param: str):
|
||||
return f"Call repetitive_call_2 tool with param {param + '_repetitive'}"
|
||||
|
||||
|
||||
def repetitive_call_2(param: str):
|
||||
return param
|
||||
|
||||
|
||||
test_case_retrieval = FilesRetrieval(
|
||||
name="test_case_retrieval",
|
||||
description="General guidance for agent test cases",
|
||||
input_dir=os.path.join(os.path.dirname(__file__), "files"),
|
||||
)
|
||||
|
||||
valid_rag_retrieval = VertexAiRagRetrieval(
|
||||
name="valid_rag_retrieval",
|
||||
rag_corpora=[
|
||||
"projects/1096655024998/locations/us-central1/ragCorpora/4985766262475849728"
|
||||
],
|
||||
description="General guidance for agent test cases",
|
||||
)
|
||||
|
||||
invalid_rag_retrieval = VertexAiRagRetrieval(
|
||||
name="invalid_rag_retrieval",
|
||||
rag_corpora=[
|
||||
"projects/1096655024998/locations/us-central1/InValidRagCorporas/4985766262475849728"
|
||||
],
|
||||
description="Invalid rag retrieval resource name",
|
||||
)
|
||||
|
||||
non_exist_rag_retrieval = VertexAiRagRetrieval(
|
||||
name="non_exist_rag_retrieval",
|
||||
rag_corpora=[
|
||||
"projects/1096655024998/locations/us-central1/RagCorpora/1234567"
|
||||
],
|
||||
description="Non exist rag retrieval resource name",
|
||||
)
|
||||
|
||||
shell_tool = LangchainTool(ShellTool())
|
||||
|
||||
docs_tool = CrewaiTool(
|
||||
name="directory_read_tool",
|
||||
description="use this to find files for you.",
|
||||
tool=DirectoryReadTool(directory="."),
|
||||
)
|
||||
|
||||
no_schema_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="no_schema_agent",
|
||||
instruction="""Just say 'Hi'
|
||||
""",
|
||||
)
|
||||
|
||||
schema_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="schema_agent",
|
||||
instruction="""
|
||||
You will be given a test case.
|
||||
Return a list of the received test case appended with '_success' and '_failure' as test_titles
|
||||
""",
|
||||
input_schema=TestCase,
|
||||
output_schema=Test,
|
||||
)
|
||||
|
||||
no_input_schema_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="no_input_schema_agent",
|
||||
instruction="""
|
||||
Just return ['Tools_success, Tools_failure']
|
||||
""",
|
||||
output_schema=Test,
|
||||
)
|
||||
|
||||
no_output_schema_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="no_output_schema_agent",
|
||||
instruction="""
|
||||
Just say 'Hi'
|
||||
""",
|
||||
input_schema=TestCase,
|
||||
)
|
||||
|
||||
single_function_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="single_function_agent",
|
||||
description="An agent that calls a single function",
|
||||
instruction="When calling tools, just return what the tool returns.",
|
||||
tools=[simple_function],
|
||||
)
|
||||
|
||||
root_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="tool_agent",
|
||||
description="An agent that can call other tools",
|
||||
instruction="When calling tools, just return what the tool returns.",
|
||||
tools=[
|
||||
simple_function,
|
||||
no_param_function,
|
||||
no_output_function,
|
||||
multiple_param_types_function,
|
||||
throw_error_function,
|
||||
list_str_param_function,
|
||||
return_list_str_function,
|
||||
# complex_function_list_dict,
|
||||
repetitive_call_1,
|
||||
repetitive_call_2,
|
||||
test_case_retrieval,
|
||||
valid_rag_retrieval,
|
||||
invalid_rag_retrieval,
|
||||
non_exist_rag_retrieval,
|
||||
shell_tool,
|
||||
docs_tool,
|
||||
AgentTool(
|
||||
agent=no_schema_agent,
|
||||
),
|
||||
AgentTool(
|
||||
agent=schema_agent,
|
||||
),
|
||||
AgentTool(
|
||||
agent=no_input_schema_agent,
|
||||
),
|
||||
AgentTool(
|
||||
agent=no_output_schema_agent,
|
||||
),
|
||||
],
|
||||
)
|
||||
Binary file not shown.
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,110 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
# https://github.com/crewAIInc/crewAI-examples/tree/main/trip_planner
|
||||
|
||||
from google.adk import Agent
|
||||
|
||||
# Agent that selects the best city for the trip.
|
||||
identify_agent = Agent(
|
||||
name='identify_agent',
|
||||
description='Select the best city based on weather, season, and prices.',
|
||||
instruction="""
|
||||
Analyze and select the best city for the trip based
|
||||
on specific criteria such as weather patterns, seasonal
|
||||
events, and travel costs. This task involves comparing
|
||||
multiple cities, considering factors like current weather
|
||||
conditions, upcoming cultural or seasonal events, and
|
||||
overall travel expenses.
|
||||
|
||||
Your final answer must be a detailed
|
||||
report on the chosen city, and everything you found out
|
||||
about it, including the actual flight costs, weather
|
||||
forecast and attractions.
|
||||
|
||||
Traveling from: {origin}
|
||||
City Options: {cities}
|
||||
Trip Date: {range}
|
||||
Traveler Interests: {interests}
|
||||
""",
|
||||
)
|
||||
|
||||
# Agent that gathers information about the city.
|
||||
gather_agent = Agent(
|
||||
name='gather_agent',
|
||||
description='Provide the BEST insights about the selected city',
|
||||
instruction="""
|
||||
As a local expert on this city you must compile an
|
||||
in-depth guide for someone traveling there and wanting
|
||||
to have THE BEST trip ever!
|
||||
Gather information about key attractions, local customs,
|
||||
special events, and daily activity recommendations.
|
||||
Find the best spots to go to, the kind of place only a
|
||||
local would know.
|
||||
This guide should provide a thorough overview of what
|
||||
the city has to offer, including hidden gems, cultural
|
||||
hotspots, must-visit landmarks, weather forecasts, and
|
||||
high level costs.
|
||||
|
||||
The final answer must be a comprehensive city guide,
|
||||
rich in cultural insights and practical tips,
|
||||
tailored to enhance the travel experience.
|
||||
|
||||
Trip Date: {range}
|
||||
Traveling from: {origin}
|
||||
Traveler Interests: {interests}
|
||||
""",
|
||||
)
|
||||
|
||||
# Agent that plans the trip.
|
||||
plan_agent = Agent(
|
||||
name='plan_agent',
|
||||
description="""Create the most amazing travel itineraries with budget and
|
||||
packing suggestions for the city""",
|
||||
instruction="""
|
||||
Expand this guide into a full 7-day travel
|
||||
itinerary with detailed per-day plans, including
|
||||
weather forecasts, places to eat, packing suggestions,
|
||||
and a budget breakdown.
|
||||
|
||||
You MUST suggest actual places to visit, actual hotels
|
||||
to stay and actual restaurants to go to.
|
||||
|
||||
This itinerary should cover all aspects of the trip,
|
||||
from arrival to departure, integrating the city guide
|
||||
information with practical travel logistics.
|
||||
|
||||
Your final answer MUST be a complete expanded travel plan,
|
||||
formatted as markdown, encompassing a daily schedule,
|
||||
anticipated weather conditions, recommended clothing and
|
||||
items to pack, and a detailed budget, ensuring THE BEST
|
||||
TRIP EVER. Be specific and give it a reason why you picked
|
||||
each place, what makes them special!
|
||||
|
||||
Trip Date: {range}
|
||||
Traveling from: {origin}
|
||||
Traveler Interests: {interests}
|
||||
""",
|
||||
)
|
||||
|
||||
root_agent = Agent(
|
||||
model='gemini-2.5-flash',
|
||||
name='trip_planner',
|
||||
description='Plan the best trip ever',
|
||||
instruction="""
|
||||
Your goal is to plan the best trip according to information listed above.
|
||||
You describe why did you choose the city, list top 3
|
||||
attractions and provide a detailed itinerary for each day.""",
|
||||
sub_agents=[identify_agent, gather_agent, plan_agent],
|
||||
)
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"criteria": {
|
||||
"response_match_score": 0.5
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"criteria": {
|
||||
"response_match_score": 0.5
|
||||
}
|
||||
}
|
||||
+64
@@ -0,0 +1,64 @@
|
||||
{
|
||||
"eval_set_id": "189d6856-9b90-4b9c-bda8-7cec899507ae",
|
||||
"name": "189d6856-9b90-4b9c-bda8-7cec899507ae",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/trip_planner_agent/test_files/trip_inquiry_sub_agent.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "1c2e8003-d19c-4912-b0ae-17b9d568f8fb",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Based on my interests, where should I go, Yosemite national park or Los Angeles?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"file_data": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"inline_data": null,
|
||||
"text": "Given your interests in food, shopping, and museums, Los Angeles would be a better choice than Yosemite National Park. Yosemite is primarily focused on outdoor activities and natural landscapes, while Los Angeles offers a diverse range of culinary experiences, shopping districts, and world-class museums. I will now gather information to create an in-depth guide for your trip to Los Angeles.\n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1747339378.484014
|
||||
}
|
||||
],
|
||||
"session_input": {
|
||||
"app_name": "trip_planner_agent",
|
||||
"user_id": "test_user",
|
||||
"state": {
|
||||
"origin": "San Francisco",
|
||||
"interests": "Food, Shopping, Museums",
|
||||
"range": "1000 miles",
|
||||
"cities": ""
|
||||
}
|
||||
},
|
||||
"creation_timestamp": 1747339378.484044
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1747339378.484056
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
{
|
||||
"eval_set_id": "e7996ccc-16bc-46bf-9a24-0a3ecc3dacd7",
|
||||
"name": "e7996ccc-16bc-46bf-9a24-0a3ecc3dacd7",
|
||||
"description": null,
|
||||
"eval_cases": [
|
||||
{
|
||||
"eval_id": "tests/integration/fixture/trip_planner_agent/trip_inquiry.test.json",
|
||||
"conversation": [
|
||||
{
|
||||
"invocation_id": "d7ff8ec1-290b-48c5-b3aa-05cb8f27b8ae",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "Hi, who are you? What can you do?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "I am trip_planner, and my goal is to plan the best trip ever. I can describe why a city was chosen, list its top attractions, and provide a detailed itinerary for each day of the trip.\n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1750190885.419684
|
||||
},
|
||||
{
|
||||
"invocation_id": "f515ff57-ff21-488f-ab92-7d7de5bb76fe",
|
||||
"user_content": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "I want to travel from San Francisco to an European country in fall next year. I am considering London and Paris. What is your advice?"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
},
|
||||
"final_response": {
|
||||
"parts": [
|
||||
{
|
||||
"video_metadata": null,
|
||||
"thought": null,
|
||||
"inline_data": null,
|
||||
"file_data": null,
|
||||
"thought_signature": null,
|
||||
"code_execution_result": null,
|
||||
"executable_code": null,
|
||||
"function_call": null,
|
||||
"function_response": null,
|
||||
"text": "Okay, I can help you analyze London and Paris to determine which city is better for your trip next fall. I will consider weather patterns, seasonal events, travel costs (including flights from San Francisco), and your interests (food, shopping, and museums). After gathering this information, I'll provide a detailed report on my chosen city.\n"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"intermediate_data": {
|
||||
"tool_uses": [
|
||||
{
|
||||
"id": null,
|
||||
"args": {
|
||||
"agent_name": "identify_agent"
|
||||
},
|
||||
"name": "transfer_to_agent"
|
||||
}
|
||||
],
|
||||
"intermediate_responses": []
|
||||
},
|
||||
"creation_timestamp": 1750190885.4197457
|
||||
}
|
||||
],
|
||||
"session_input": {
|
||||
"app_name": "trip_planner_agent",
|
||||
"user_id": "test_user",
|
||||
"state": {
|
||||
"origin": "San Francisco",
|
||||
"interests": "Food, Shopping, Museums",
|
||||
"range": "1000 miles",
|
||||
"cities": ""
|
||||
}
|
||||
},
|
||||
"creation_timestamp": 1750190885.4197533
|
||||
}
|
||||
],
|
||||
"creation_timestamp": 1750190885.4197605
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
# Integration tests for GCP Agent Identity Credentials service
|
||||
|
||||
Verifies OAuth flows using GCP Agent Identity Credentials service.
|
||||
|
||||
## Setup
|
||||
|
||||
To set up your environment for the first time, create a virtual environment
|
||||
and install dependencies:
|
||||
```bash
|
||||
uv venv --python "python3.11" ".venv"
|
||||
source .venv/bin/activate
|
||||
uv sync --all-extras
|
||||
```
|
||||
|
||||
Then, install test specific packages
|
||||
```bash
|
||||
pip install google-cloud-iamconnectorcredentials
|
||||
```
|
||||
|
||||
## Run Tests
|
||||
```bash
|
||||
pytest -s tests/integration/integrations/agent_identity
|
||||
```
|
||||
@@ -0,0 +1,268 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""E2E Integration Test for GCP Agent Identity Auth Provider two-legged OAuth Flow."""
|
||||
|
||||
import dataclasses
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
|
||||
from google.adk import Agent
|
||||
from google.adk import Runner
|
||||
from google.adk.auth.auth_tool import AuthConfig
|
||||
from google.adk.auth.credential_manager import CredentialManager
|
||||
from google.adk.integrations.agent_identity import _iam_connector_credentials_provider
|
||||
from google.adk.integrations.agent_identity import GcpAuthProvider
|
||||
from google.adk.integrations.agent_identity import GcpAuthProviderScheme
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.tools.base_authenticated_tool import BaseAuthenticatedTool
|
||||
from google.adk.tools.mcp_tool.mcp_tool import McpTool
|
||||
from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsRequest
|
||||
from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsResponse
|
||||
from google.genai import types
|
||||
from mcp.types import Tool as McpBaseTool
|
||||
import pytest
|
||||
|
||||
from tests.unittests import testing_utils
|
||||
|
||||
DUMMY_TOKEN = "fake-gcp-2lo-token-123"
|
||||
TEST_CONNECTOR_2LO = (
|
||||
"projects/test-project/locations/global/connectors/test-connector"
|
||||
)
|
||||
|
||||
|
||||
class DummyTool(BaseAuthenticatedTool):
|
||||
|
||||
def __init__(self, auth_config: AuthConfig) -> None:
|
||||
super().__init__(
|
||||
name="dummy_tool",
|
||||
description="Dummy tool for testing 2LO.",
|
||||
auth_config=auth_config,
|
||||
)
|
||||
|
||||
def _get_declaration(self) -> types.FunctionDeclaration:
|
||||
return types.FunctionDeclaration(
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
parameters=types.Schema(
|
||||
type="OBJECT",
|
||||
properties={},
|
||||
),
|
||||
)
|
||||
|
||||
async def _run_async_impl(
|
||||
self, *, args: dict[str, Any] | None, tool_context: Any, credential: Any
|
||||
) -> Any:
|
||||
# Return the token to prove the provider gave the expected credential
|
||||
if credential.http and credential.http.credentials:
|
||||
return credential.http.credentials.token
|
||||
if credential.oauth2 and credential.oauth2.access_token:
|
||||
return credential.oauth2.access_token
|
||||
return None
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _DummyOperation:
|
||||
done: bool = True
|
||||
error: Any = None
|
||||
metadata: Any = None
|
||||
response: Any = dataclasses.field(init=False)
|
||||
operation: Any = dataclasses.field(init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.response = mock.Mock()
|
||||
mock_credential = RetrieveCredentialsResponse(
|
||||
header="Authorization: Bearer", token=DUMMY_TOKEN
|
||||
)
|
||||
self.response.value = RetrieveCredentialsResponse.serialize(mock_credential)
|
||||
self.operation = self
|
||||
|
||||
def HasField(self, field_name: str) -> bool:
|
||||
return getattr(self, field_name, None) is not None
|
||||
|
||||
|
||||
# Mocked execution; pin to a single LLM backend to avoid duplicate runs.
|
||||
@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"], indirect=True)
|
||||
@pytest.mark.asyncio
|
||||
async def test_gcp_agent_identity_2lo_gets_token() -> None:
|
||||
"""Test the end-to-end flow fetching 2LO OAuth token from GCP Agent Identity credentials service."""
|
||||
|
||||
# Clear registry to isolate tests
|
||||
CredentialManager._auth_provider_registry._providers.clear()
|
||||
|
||||
# 1. Setup mocked GCP Client to return the fake Bearer token
|
||||
with mock.patch.object(
|
||||
_iam_connector_credentials_provider,
|
||||
"Client",
|
||||
autospec=True,
|
||||
) as mock_client_cls:
|
||||
|
||||
mock_operation = _DummyOperation()
|
||||
|
||||
mock_client_cls.return_value.retrieve_credentials.return_value = (
|
||||
mock_operation
|
||||
)
|
||||
|
||||
# 2. Configure Auth and DummyTool
|
||||
auth_scheme = GcpAuthProviderScheme(
|
||||
name=TEST_CONNECTOR_2LO,
|
||||
scopes=["test-scope"],
|
||||
)
|
||||
auth_config = AuthConfig(auth_scheme=auth_scheme)
|
||||
dummy_tool = DummyTool(auth_config=auth_config)
|
||||
|
||||
# 3. Setup LLM, Agent, and Runner
|
||||
# We mock the LLM to just issue the tool call to 'dummy_tool'
|
||||
mock_model = testing_utils.MockModel.create(
|
||||
responses=[
|
||||
types.Part.from_function_call(name="dummy_tool", args={}),
|
||||
"Tool executed successfully.",
|
||||
]
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
name="test_agent",
|
||||
model=mock_model,
|
||||
instruction="You are an agent. Use the dummy_tool when needed.",
|
||||
tools=[dummy_tool],
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app_name="test_mcp_2lo_app",
|
||||
agent=agent,
|
||||
session_service=InMemorySessionService(),
|
||||
auto_create_session=True,
|
||||
)
|
||||
|
||||
# 4. Register Auth Provider
|
||||
CredentialManager.register_auth_provider(GcpAuthProvider())
|
||||
|
||||
# 5. Execute Flow
|
||||
event_list = []
|
||||
async for event in runner.run_async(
|
||||
user_id="test_user",
|
||||
session_id="test_session1",
|
||||
new_message=types.UserContent(
|
||||
parts=[types.Part(text="Get me the token.")]
|
||||
),
|
||||
):
|
||||
event_list.append(event)
|
||||
|
||||
# 6. Assertions
|
||||
|
||||
# Assert GCP Agent Identity client was invoked for credentials
|
||||
expected_request = RetrieveCredentialsRequest(
|
||||
connector=TEST_CONNECTOR_2LO,
|
||||
user_id="test_user",
|
||||
scopes=["test-scope"],
|
||||
continue_uri="",
|
||||
force_refresh=False,
|
||||
)
|
||||
mock_client_cls.return_value.retrieve_credentials.assert_called_once_with(
|
||||
expected_request
|
||||
)
|
||||
|
||||
# 3 Events: Model FunctionCall -> Tool FunctionResponse -> Final LLM Text
|
||||
assert len(event_list) == 3
|
||||
last_event = event_list[-1]
|
||||
assert last_event.content.parts[0].text == "Tool executed successfully."
|
||||
|
||||
# Validate that the mock model received the query and the tool callback
|
||||
requests = mock_model.requests
|
||||
# 2 Events: User Input -> Tool FunctionResponse
|
||||
assert len(requests) == 2
|
||||
|
||||
# Extract the function response from the prompt payload sent to the LLM
|
||||
last_request = requests[-1]
|
||||
function_response = next(
|
||||
(
|
||||
p.function_response
|
||||
for p in last_request.contents[-1].parts
|
||||
if p.function_response
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
assert function_response.name == "dummy_tool"
|
||||
assert DUMMY_TOKEN in str(function_response.response)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"], indirect=True)
|
||||
@pytest.mark.asyncio
|
||||
async def test_gcp_agent_identity_2lo_sends_authorization_header_to_mcp_session(
|
||||
llm_backend: Any,
|
||||
) -> None:
|
||||
"""Ensures a 2LO token from GCP is correctly passed into the outbound MCP session headers."""
|
||||
CredentialManager._auth_provider_registry._providers.clear()
|
||||
CredentialManager.register_auth_provider(GcpAuthProvider())
|
||||
|
||||
mock_operation = _DummyOperation()
|
||||
with mock.patch.object(
|
||||
_iam_connector_credentials_provider, "Client", autospec=True
|
||||
) as mock_gcp:
|
||||
mock_gcp.return_value.retrieve_credentials.return_value = mock_operation
|
||||
|
||||
mock_session_mgr = mock.AsyncMock()
|
||||
mock_session_mgr.create_session.return_value.call_tool.return_value = (
|
||||
mock.MagicMock()
|
||||
)
|
||||
|
||||
mcp_tool = McpTool(
|
||||
mcp_tool=McpBaseTool(
|
||||
name="dummy_mcp",
|
||||
description="Dummy MCP tool for testing.",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
),
|
||||
mcp_session_manager=mock_session_mgr,
|
||||
auth_scheme=GcpAuthProviderScheme(
|
||||
name=TEST_CONNECTOR_2LO, scopes=["test-scope"]
|
||||
),
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
name="test_agent",
|
||||
model=testing_utils.MockModel.create(
|
||||
responses=[
|
||||
types.Part.from_function_call(name="dummy_mcp", args={}),
|
||||
"Tool executed successfully.",
|
||||
]
|
||||
),
|
||||
instruction="Use dummy_mcp tool.",
|
||||
tools=[mcp_tool],
|
||||
)
|
||||
|
||||
async for _ in Runner(
|
||||
app_name="test_mcp_header_app",
|
||||
agent=agent,
|
||||
session_service=InMemorySessionService(),
|
||||
auto_create_session=True,
|
||||
).run_async(
|
||||
user_id="test_user",
|
||||
session_id="session-id-2",
|
||||
new_message=types.UserContent(parts=[types.Part(text="Run tool.")]),
|
||||
):
|
||||
pass
|
||||
|
||||
mock_gcp.return_value.retrieve_credentials.assert_called_once_with(
|
||||
RetrieveCredentialsRequest(
|
||||
connector=TEST_CONNECTOR_2LO,
|
||||
user_id="test_user",
|
||||
scopes=["test-scope"],
|
||||
force_refresh=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert mock_session_mgr.create_session.call_args.kwargs.get("headers") == {
|
||||
"Authorization": f"Bearer {DUMMY_TOKEN}"
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""E2E Integration Test for 3LO flow using GCP Agent Identity service."""
|
||||
|
||||
import dataclasses
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
|
||||
from google.adk import Agent
|
||||
from google.adk import Runner
|
||||
from google.adk.auth.auth_tool import AuthConfig
|
||||
from google.adk.auth.credential_manager import CredentialManager
|
||||
from google.adk.integrations.agent_identity import _iam_connector_credentials_provider
|
||||
from google.adk.integrations.agent_identity import GcpAuthProvider
|
||||
from google.adk.integrations.agent_identity import GcpAuthProviderScheme
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.tools.base_authenticated_tool import BaseAuthenticatedTool
|
||||
from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsMetadata
|
||||
from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsRequest
|
||||
from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsResponse
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
from tests.unittests import testing_utils
|
||||
|
||||
DUMMY_TOKEN = "mock-token-3legged"
|
||||
TEST_CONNECTOR_3LO = (
|
||||
"projects/my-project/locations/some-location/connectors/test-connector-3lo"
|
||||
)
|
||||
|
||||
|
||||
class DummyTool(BaseAuthenticatedTool):
|
||||
|
||||
def __init__(self, auth_config: AuthConfig) -> None:
|
||||
super().__init__(
|
||||
name="dummy_tool",
|
||||
description="Dummy tool for testing 3LO.",
|
||||
auth_config=auth_config,
|
||||
)
|
||||
|
||||
def _get_declaration(self) -> types.FunctionDeclaration:
|
||||
return types.FunctionDeclaration(
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
parameters=types.Schema(
|
||||
type="OBJECT",
|
||||
properties={},
|
||||
),
|
||||
)
|
||||
|
||||
async def _run_async_impl(
|
||||
self, *, args: dict[str, Any] | None, tool_context: Any, credential: Any
|
||||
) -> Any:
|
||||
# Extract and return the token to prove the provider gave us the expected credential
|
||||
if credential.http and credential.http.credentials:
|
||||
return credential.http.credentials.token
|
||||
if credential.oauth2 and credential.oauth2.access_token:
|
||||
return credential.oauth2.access_token
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _MockOperation:
|
||||
done: bool
|
||||
response_obj: Any = None
|
||||
metadata_obj: Any = None
|
||||
error: Any = None
|
||||
metadata: Any = dataclasses.field(init=False, default=None)
|
||||
response: Any = dataclasses.field(init=False, default=None)
|
||||
operation: Any = dataclasses.field(init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.metadata_obj:
|
||||
self.metadata = mock.Mock()
|
||||
self.metadata.value = RetrieveCredentialsMetadata.serialize(
|
||||
self.metadata_obj
|
||||
)
|
||||
if self.response_obj:
|
||||
self.response = mock.Mock()
|
||||
self.response.value = RetrieveCredentialsResponse.serialize(
|
||||
self.response_obj
|
||||
)
|
||||
self.operation = self
|
||||
|
||||
def HasField(self, field_name: str) -> bool:
|
||||
return getattr(self, field_name, None) is not None
|
||||
|
||||
|
||||
class MockGcpClient:
|
||||
"""Lightweight in-memory mock for Agent Identity Credentials service 3LO Consent Flow."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.finalized_connectors = set()
|
||||
|
||||
def retrieve_credentials(
|
||||
self,
|
||||
request: RetrieveCredentialsRequest | dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> _MockOperation:
|
||||
connector = (
|
||||
request.get("connector")
|
||||
if isinstance(request, dict)
|
||||
else getattr(request, "connector", None)
|
||||
)
|
||||
|
||||
if connector in self.finalized_connectors:
|
||||
mock_credential = RetrieveCredentialsResponse(
|
||||
token=DUMMY_TOKEN, header="Authorization: Bearer"
|
||||
)
|
||||
return _MockOperation(done=True, response_obj=mock_credential)
|
||||
|
||||
# Otherwise, return Consent Required
|
||||
# Auto-finalize for the next call to simulate user approval flow
|
||||
self.finalized_connectors.add(connector)
|
||||
|
||||
mock_metadata = RetrieveCredentialsMetadata(
|
||||
uri_consent_required=RetrieveCredentialsMetadata.UriConsentRequired(
|
||||
authorization_uri="http://mock-auth-uri",
|
||||
consent_nonce="mock-consent-nonce",
|
||||
)
|
||||
)
|
||||
return _MockOperation(done=False, metadata_obj=mock_metadata)
|
||||
|
||||
|
||||
# Mocked execution; pin to a single LLM backend to avoid duplicate runs.
|
||||
@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"], indirect=True)
|
||||
@pytest.mark.asyncio
|
||||
async def test_gcp_agent_identity_3lo_user_consent_flow() -> None:
|
||||
# Clear registry to isolate tests
|
||||
CredentialManager._auth_provider_registry._providers.clear()
|
||||
|
||||
# 1. Setup mocked GCP Client to simulate stateful 3LO process
|
||||
mock_gcp_client = MockGcpClient()
|
||||
|
||||
with mock.patch.object(
|
||||
_iam_connector_credentials_provider,
|
||||
"Client",
|
||||
autospec=True,
|
||||
) as mock_client_cls:
|
||||
mock_client_cls.return_value.retrieve_credentials.side_effect = (
|
||||
mock_gcp_client.retrieve_credentials
|
||||
)
|
||||
|
||||
# 2. Configure Auth and DummyTool
|
||||
auth_scheme = GcpAuthProviderScheme(
|
||||
name=TEST_CONNECTOR_3LO,
|
||||
scopes=["test-scope"],
|
||||
continue_uri="https://example.com/continue",
|
||||
)
|
||||
auth_config = AuthConfig(auth_scheme=auth_scheme)
|
||||
dummy_tool = DummyTool(auth_config=auth_config)
|
||||
|
||||
# 3. Setup LLM, Agent, and Runner
|
||||
# We mock the LLM to just issue the tool call to 'dummy_tool'
|
||||
mock_model = testing_utils.MockModel.create(
|
||||
responses=[
|
||||
types.Part.from_function_call(name="dummy_tool", args={}),
|
||||
"Tool executed successfully.",
|
||||
]
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
name="test_agent",
|
||||
model=mock_model,
|
||||
instruction="You are an agent. Use the dummy_tool when needed.",
|
||||
tools=[dummy_tool],
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app_name="test_mcp_3lo_app",
|
||||
agent=agent,
|
||||
session_service=InMemorySessionService(),
|
||||
auto_create_session=True,
|
||||
)
|
||||
|
||||
# 4. Register Auth Provider
|
||||
CredentialManager.register_auth_provider(GcpAuthProvider())
|
||||
|
||||
# 5. Execute Flow
|
||||
session = await runner.session_service.create_session(
|
||||
app_name="test_mcp_3lo_app", user_id="test_user"
|
||||
)
|
||||
|
||||
event_list = []
|
||||
|
||||
# Step 5a: User sends message, Agent requests credential
|
||||
async for event in runner.run_async(
|
||||
user_id="test_user",
|
||||
session_id=session.id,
|
||||
new_message=types.UserContent(
|
||||
parts=[types.Part(text="Get me the token.")]
|
||||
),
|
||||
):
|
||||
event_list.append(event)
|
||||
|
||||
def _find_auth_request_event(events):
|
||||
for event in events:
|
||||
for part in event.content.parts:
|
||||
if (
|
||||
part.function_call
|
||||
and part.function_call.name == "adk_request_credential"
|
||||
):
|
||||
return event
|
||||
return None
|
||||
|
||||
auth_request_event = _find_auth_request_event(event_list)
|
||||
|
||||
assert (
|
||||
auth_request_event
|
||||
), "Expected adk_request_credential tool call not found."
|
||||
|
||||
# Step 5b: Simulate User Consent
|
||||
call_part = next(
|
||||
p for p in auth_request_event.content.parts if p.function_call
|
||||
)
|
||||
request_auth_config = call_part.function_call.args.get("authConfig", {})
|
||||
|
||||
assert (
|
||||
request_auth_config.get("exchangedAuthCredential", {})
|
||||
.get("oauth2", {})
|
||||
.get("nonce")
|
||||
== "mock-consent-nonce"
|
||||
)
|
||||
|
||||
# Step 5c: User acknowledges credential request
|
||||
response_part = types.Part.from_function_response(
|
||||
name="adk_request_credential", response=request_auth_config
|
||||
)
|
||||
response_part.function_response.id = call_part.function_call.id
|
||||
|
||||
final_response_parts = []
|
||||
async for event in runner.run_async(
|
||||
user_id="test_user",
|
||||
session_id=session.id,
|
||||
new_message=types.UserContent(parts=[response_part]),
|
||||
):
|
||||
event_list.append(event)
|
||||
if event.content:
|
||||
for part in event.content.parts:
|
||||
if part.text:
|
||||
final_response_parts.append(part.text)
|
||||
|
||||
final_response_text = "".join(final_response_parts)
|
||||
|
||||
# 6. Assertions
|
||||
|
||||
# Assert GCP Agent Identity client was invoked for credentials twice
|
||||
# (Initial Request + Post-Consent call)
|
||||
assert mock_client_cls.return_value.retrieve_credentials.call_count == 2
|
||||
expected_request = RetrieveCredentialsRequest(
|
||||
connector=TEST_CONNECTOR_3LO,
|
||||
user_id="test_user",
|
||||
scopes=["test-scope"],
|
||||
continue_uri="https://example.com/continue",
|
||||
force_refresh=False,
|
||||
)
|
||||
mock_client_cls.return_value.retrieve_credentials.assert_called_with(
|
||||
expected_request
|
||||
)
|
||||
|
||||
assert "Tool executed successfully." in final_response_text
|
||||
|
||||
# Validate requests received by the mock model
|
||||
requests = mock_model.requests
|
||||
# Two LLM calls: the tool call in turn 1 (which ends at the EUC) and the
|
||||
# post-consent response in turn 2.
|
||||
assert len(requests) == 2
|
||||
|
||||
# Extract the function response from the prompt payload sent to the LLM
|
||||
last_request = requests[-1]
|
||||
function_response = next(
|
||||
(
|
||||
p.function_response
|
||||
for p in last_request.contents[-1].parts
|
||||
if p.function_response
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
assert function_response is not None
|
||||
assert function_response.name == "dummy_tool"
|
||||
assert DUMMY_TOKEN in str(function_response.response)
|
||||
@@ -0,0 +1,263 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""E2E Integration Test for GCP Agent Identity Auth Provider two-legged OAuth Flow."""
|
||||
|
||||
import dataclasses
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
|
||||
from google.adk import Agent
|
||||
from google.adk import Runner
|
||||
from google.adk.auth.auth_tool import AuthConfig
|
||||
from google.adk.auth.credential_manager import CredentialManager
|
||||
from google.adk.integrations.agent_identity import _agent_identity_credentials_provider
|
||||
from google.adk.integrations.agent_identity import GcpAuthProvider
|
||||
from google.adk.integrations.agent_identity import GcpAuthProviderScheme
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.tools.base_authenticated_tool import BaseAuthenticatedTool
|
||||
from google.adk.tools.mcp_tool.mcp_tool import McpTool
|
||||
from google.cloud.agentidentitycredentials_v1 import RetrieveCredentialsRequest
|
||||
from google.genai import types
|
||||
from mcp.types import Tool as McpBaseTool
|
||||
import pytest
|
||||
|
||||
from tests.unittests import testing_utils
|
||||
|
||||
DUMMY_TOKEN = "fake-gcp-2lo-token-123"
|
||||
TEST_AUTH_PROVIDER_2LO = (
|
||||
"projects/test-project/locations/global/authProviders/test-provider"
|
||||
)
|
||||
|
||||
|
||||
class DummyTool(BaseAuthenticatedTool):
|
||||
|
||||
def __init__(self, auth_config: AuthConfig) -> None:
|
||||
super().__init__(
|
||||
name="dummy_tool",
|
||||
description="Dummy tool for testing 2LO.",
|
||||
auth_config=auth_config,
|
||||
)
|
||||
|
||||
def _get_declaration(self) -> types.FunctionDeclaration:
|
||||
return types.FunctionDeclaration(
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
parameters=types.Schema(
|
||||
type="OBJECT",
|
||||
properties={},
|
||||
),
|
||||
)
|
||||
|
||||
async def _run_async_impl(
|
||||
self, *, args: dict[str, Any] | None, tool_context: Any, credential: Any
|
||||
) -> Any:
|
||||
# Return the token to prove the provider gave the expected credential
|
||||
if credential.http and credential.http.credentials:
|
||||
return credential.http.credentials.token
|
||||
if credential.oauth2 and credential.oauth2.access_token:
|
||||
return credential.oauth2.access_token
|
||||
return None
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _DummyOperation:
|
||||
done: bool = True
|
||||
error: Any = None
|
||||
metadata: Any = None
|
||||
response: Any = dataclasses.field(init=False)
|
||||
success: Any = dataclasses.field(init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.success = mock.Mock()
|
||||
self.success.header = "Authorization: Bearer"
|
||||
self.success.token = DUMMY_TOKEN
|
||||
self.response = self
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
return key == "success"
|
||||
|
||||
|
||||
# Mocked execution; pin to a single LLM backend to avoid duplicate runs.
|
||||
@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"], indirect=True)
|
||||
@pytest.mark.asyncio
|
||||
async def test_gcp_agent_identity_2lo_gets_token() -> None:
|
||||
"""Test the end-to-end flow fetching 2LO OAuth token from GCP Agent Identity credentials service."""
|
||||
|
||||
# Clear registry to isolate tests
|
||||
CredentialManager._auth_provider_registry._providers.clear()
|
||||
|
||||
# 1. Setup mocked GCP Client to return the fake Bearer token
|
||||
with mock.patch.object(
|
||||
_agent_identity_credentials_provider,
|
||||
"Client",
|
||||
autospec=True,
|
||||
) as mock_client_cls:
|
||||
|
||||
mock_operation = _DummyOperation()
|
||||
|
||||
mock_client_cls.return_value.retrieve_credentials.return_value = (
|
||||
mock_operation
|
||||
)
|
||||
|
||||
# 2. Configure Auth and DummyTool
|
||||
auth_scheme = GcpAuthProviderScheme(
|
||||
name=TEST_AUTH_PROVIDER_2LO,
|
||||
scopes=["test-scope"],
|
||||
)
|
||||
auth_config = AuthConfig(auth_scheme=auth_scheme)
|
||||
dummy_tool = DummyTool(auth_config=auth_config)
|
||||
|
||||
# 3. Setup LLM, Agent, and Runner
|
||||
# We mock the LLM to just issue the tool call to 'dummy_tool'
|
||||
mock_model = testing_utils.MockModel.create(
|
||||
responses=[
|
||||
types.Part.from_function_call(name="dummy_tool", args={}),
|
||||
"Tool executed successfully.",
|
||||
]
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
name="test_agent",
|
||||
model=mock_model,
|
||||
instruction="You are an agent. Use the dummy_tool when needed.",
|
||||
tools=[dummy_tool],
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app_name="test_mcp_2lo_app",
|
||||
agent=agent,
|
||||
session_service=InMemorySessionService(),
|
||||
auto_create_session=True,
|
||||
)
|
||||
|
||||
# 4. Register Auth Provider
|
||||
CredentialManager.register_auth_provider(GcpAuthProvider())
|
||||
|
||||
# 5. Execute Flow
|
||||
event_list = []
|
||||
async for event in runner.run_async(
|
||||
user_id="test_user",
|
||||
session_id="test_session1",
|
||||
new_message=types.UserContent(
|
||||
parts=[types.Part(text="Get me the token.")]
|
||||
),
|
||||
):
|
||||
event_list.append(event)
|
||||
|
||||
# 6. Assertions
|
||||
|
||||
# Assert GCP Agent Identity client was invoked for credentials
|
||||
expected_request = RetrieveCredentialsRequest(
|
||||
auth_provider=TEST_AUTH_PROVIDER_2LO,
|
||||
user_id="test_user",
|
||||
scopes=["test-scope"],
|
||||
continue_uri="",
|
||||
)
|
||||
mock_client_cls.return_value.retrieve_credentials.assert_called_once_with(
|
||||
expected_request
|
||||
)
|
||||
|
||||
# 3 Events: Model FunctionCall -> Tool FunctionResponse -> Final LLM Text
|
||||
assert len(event_list) == 3
|
||||
last_event = event_list[-1]
|
||||
assert last_event.content.parts[0].text == "Tool executed successfully."
|
||||
|
||||
# Validate that the mock model received the query and the tool callback
|
||||
requests = mock_model.requests
|
||||
# 2 Events: User Input -> Tool FunctionResponse
|
||||
assert len(requests) == 2
|
||||
|
||||
# Extract the function response from the prompt payload sent to the LLM
|
||||
last_request = requests[-1]
|
||||
function_response = next(
|
||||
(
|
||||
p.function_response
|
||||
for p in last_request.contents[-1].parts
|
||||
if p.function_response
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
assert function_response.name == "dummy_tool"
|
||||
assert DUMMY_TOKEN in str(function_response.response)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"], indirect=True)
|
||||
@pytest.mark.asyncio
|
||||
async def test_gcp_agent_identity_2lo_sends_authorization_header_to_mcp_session(
|
||||
llm_backend: Any,
|
||||
) -> None:
|
||||
"""Ensures a 2LO token from GCP is correctly passed into the outbound MCP session headers."""
|
||||
CredentialManager._auth_provider_registry._providers.clear()
|
||||
CredentialManager.register_auth_provider(GcpAuthProvider())
|
||||
|
||||
mock_operation = _DummyOperation()
|
||||
with mock.patch.object(
|
||||
_agent_identity_credentials_provider, "Client", autospec=True
|
||||
) as mock_gcp:
|
||||
mock_gcp.return_value.retrieve_credentials.return_value = mock_operation
|
||||
|
||||
mock_session_mgr = mock.AsyncMock()
|
||||
mock_session_mgr.create_session.return_value.call_tool.return_value = (
|
||||
mock.MagicMock()
|
||||
)
|
||||
|
||||
mcp_tool = McpTool(
|
||||
mcp_tool=McpBaseTool(
|
||||
name="dummy_mcp",
|
||||
description="Dummy MCP tool for testing.",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
),
|
||||
mcp_session_manager=mock_session_mgr,
|
||||
auth_scheme=GcpAuthProviderScheme(
|
||||
name=TEST_AUTH_PROVIDER_2LO, scopes=["test-scope"]
|
||||
),
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
name="test_agent",
|
||||
model=testing_utils.MockModel.create(
|
||||
responses=[
|
||||
types.Part.from_function_call(name="dummy_mcp", args={}),
|
||||
"Tool executed successfully.",
|
||||
]
|
||||
),
|
||||
instruction="Use dummy_mcp tool.",
|
||||
tools=[mcp_tool],
|
||||
)
|
||||
|
||||
async for _ in Runner(
|
||||
app_name="test_mcp_header_app",
|
||||
agent=agent,
|
||||
session_service=InMemorySessionService(),
|
||||
auto_create_session=True,
|
||||
).run_async(
|
||||
user_id="test_user",
|
||||
session_id="session-id-2",
|
||||
new_message=types.UserContent(parts=[types.Part(text="Run tool.")]),
|
||||
):
|
||||
pass
|
||||
|
||||
mock_gcp.return_value.retrieve_credentials.assert_called_once_with(
|
||||
RetrieveCredentialsRequest(
|
||||
auth_provider=TEST_AUTH_PROVIDER_2LO,
|
||||
user_id="test_user",
|
||||
scopes=["test-scope"],
|
||||
)
|
||||
)
|
||||
|
||||
assert mock_session_mgr.create_session.call_args.kwargs.get("headers") == {
|
||||
"Authorization": f"Bearer {DUMMY_TOKEN}"
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""E2E Integration Test for 3LO flow using GCP Agent Identity service."""
|
||||
|
||||
import dataclasses
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
|
||||
from google.adk import Agent
|
||||
from google.adk import Runner
|
||||
from google.adk.auth.auth_tool import AuthConfig
|
||||
from google.adk.auth.credential_manager import CredentialManager
|
||||
from google.adk.integrations.agent_identity import _agent_identity_credentials_provider
|
||||
from google.adk.integrations.agent_identity import GcpAuthProvider
|
||||
from google.adk.integrations.agent_identity import GcpAuthProviderScheme
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.tools.base_authenticated_tool import BaseAuthenticatedTool
|
||||
from google.cloud.agentidentitycredentials_v1 import RetrieveCredentialsRequest
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
from tests.unittests import testing_utils
|
||||
|
||||
DUMMY_TOKEN = "mock-token-3legged"
|
||||
TEST_AUTH_PROVIDER_3LO = "projects/my-project/locations/some-location/authProviders/test-provider-3lo"
|
||||
|
||||
|
||||
class DummyTool(BaseAuthenticatedTool):
|
||||
|
||||
def __init__(self, auth_config: AuthConfig) -> None:
|
||||
super().__init__(
|
||||
name="dummy_tool",
|
||||
description="Dummy tool for testing 3LO.",
|
||||
auth_config=auth_config,
|
||||
)
|
||||
|
||||
def _get_declaration(self) -> types.FunctionDeclaration:
|
||||
return types.FunctionDeclaration(
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
parameters=types.Schema(
|
||||
type="OBJECT",
|
||||
properties={},
|
||||
),
|
||||
)
|
||||
|
||||
async def _run_async_impl(
|
||||
self, *, args: dict[str, Any] | None, tool_context: Any, credential: Any
|
||||
) -> Any:
|
||||
# Extract and return the token to prove the provider gave us the expected credential
|
||||
if credential.http and credential.http.credentials:
|
||||
return credential.http.credentials.token
|
||||
if credential.oauth2 and credential.oauth2.access_token:
|
||||
return credential.oauth2.access_token
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _MockOperation:
|
||||
done: bool
|
||||
response_obj: Any = None
|
||||
metadata_obj: Any = None
|
||||
error: Any = None
|
||||
response: Any = dataclasses.field(init=False, default=None)
|
||||
success: Any = dataclasses.field(init=False, default=None)
|
||||
uri_consent_required: Any = dataclasses.field(init=False, default=None)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.metadata_obj:
|
||||
self.uri_consent_required = self.metadata_obj.uri_consent_required
|
||||
if self.response_obj:
|
||||
self.success = mock.Mock()
|
||||
self.success.header = self.response_obj.header
|
||||
self.success.token = self.response_obj.token
|
||||
self.response = self
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
return getattr(self, key, None) is not None
|
||||
|
||||
|
||||
class MockGcpClient:
|
||||
"""Lightweight in-memory mock for Agent Identity Credentials service 3LO Consent Flow."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.finalized_connectors = set()
|
||||
|
||||
def retrieve_credentials(
|
||||
self,
|
||||
request: RetrieveCredentialsRequest | dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> _MockOperation:
|
||||
auth_provider = (
|
||||
request.get("auth_provider")
|
||||
if isinstance(request, dict)
|
||||
else getattr(request, "auth_provider", None)
|
||||
)
|
||||
|
||||
if auth_provider in self.finalized_connectors:
|
||||
mock_credential = mock.Mock(
|
||||
token=DUMMY_TOKEN, header="Authorization: Bearer"
|
||||
)
|
||||
return _MockOperation(done=True, response_obj=mock_credential)
|
||||
|
||||
# Otherwise, return Consent Required
|
||||
# Auto-finalize for the next call to simulate user approval flow
|
||||
self.finalized_connectors.add(auth_provider)
|
||||
|
||||
mock_metadata = mock.Mock()
|
||||
mock_metadata.uri_consent_required = mock.Mock(
|
||||
authorization_uri="http://mock-auth-uri",
|
||||
consent_nonce="mock-consent-nonce",
|
||||
)
|
||||
return _MockOperation(done=False, metadata_obj=mock_metadata)
|
||||
|
||||
|
||||
# Mocked execution; pin to a single LLM backend to avoid duplicate runs.
|
||||
@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"], indirect=True)
|
||||
@pytest.mark.asyncio
|
||||
async def test_gcp_agent_identity_3lo_user_consent_flow() -> None:
|
||||
# Clear registry to isolate tests
|
||||
CredentialManager._auth_provider_registry._providers.clear()
|
||||
|
||||
# 1. Setup mocked GCP Client to simulate stateful 3LO process
|
||||
mock_gcp_client = MockGcpClient()
|
||||
|
||||
with mock.patch.object(
|
||||
_agent_identity_credentials_provider,
|
||||
"Client",
|
||||
autospec=True,
|
||||
) as mock_client_cls:
|
||||
mock_client_cls.return_value.retrieve_credentials.side_effect = (
|
||||
mock_gcp_client.retrieve_credentials
|
||||
)
|
||||
|
||||
# 2. Configure Auth and DummyTool
|
||||
auth_scheme = GcpAuthProviderScheme(
|
||||
name=TEST_AUTH_PROVIDER_3LO,
|
||||
scopes=["test-scope"],
|
||||
continue_uri="https://example.com/continue",
|
||||
)
|
||||
auth_config = AuthConfig(auth_scheme=auth_scheme)
|
||||
dummy_tool = DummyTool(auth_config=auth_config)
|
||||
|
||||
# 3. Setup LLM, Agent, and Runner
|
||||
# We mock the LLM to just issue the tool call to 'dummy_tool'
|
||||
mock_model = testing_utils.MockModel.create(
|
||||
responses=[
|
||||
types.Part.from_function_call(name="dummy_tool", args={}),
|
||||
"Tool executed successfully.",
|
||||
]
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
name="test_agent",
|
||||
model=mock_model,
|
||||
instruction="You are an agent. Use the dummy_tool when needed.",
|
||||
tools=[dummy_tool],
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app_name="test_mcp_3lo_app",
|
||||
agent=agent,
|
||||
session_service=InMemorySessionService(),
|
||||
auto_create_session=True,
|
||||
)
|
||||
|
||||
# 4. Register Auth Provider
|
||||
CredentialManager.register_auth_provider(GcpAuthProvider())
|
||||
|
||||
# 5. Execute Flow
|
||||
session = await runner.session_service.create_session(
|
||||
app_name="test_mcp_3lo_app", user_id="test_user"
|
||||
)
|
||||
|
||||
event_list = []
|
||||
|
||||
# Step 5a: User sends message, Agent requests credential
|
||||
async for event in runner.run_async(
|
||||
user_id="test_user",
|
||||
session_id=session.id,
|
||||
new_message=types.UserContent(
|
||||
parts=[types.Part(text="Get me the token.")]
|
||||
),
|
||||
):
|
||||
event_list.append(event)
|
||||
|
||||
def _find_auth_request_event(events):
|
||||
for event in events:
|
||||
for part in event.content.parts:
|
||||
if (
|
||||
part.function_call
|
||||
and part.function_call.name == "adk_request_credential"
|
||||
):
|
||||
return event
|
||||
return None
|
||||
|
||||
auth_request_event = _find_auth_request_event(event_list)
|
||||
|
||||
assert (
|
||||
auth_request_event
|
||||
), "Expected adk_request_credential tool call not found."
|
||||
|
||||
# Step 5b: Simulate User Consent
|
||||
call_part = next(
|
||||
p for p in auth_request_event.content.parts if p.function_call
|
||||
)
|
||||
request_auth_config = call_part.function_call.args.get("authConfig", {})
|
||||
|
||||
assert (
|
||||
request_auth_config.get("exchangedAuthCredential", {})
|
||||
.get("oauth2", {})
|
||||
.get("nonce")
|
||||
== "mock-consent-nonce"
|
||||
)
|
||||
|
||||
# Step 5c: User acknowledges credential request
|
||||
response_part = types.Part.from_function_response(
|
||||
name="adk_request_credential", response=request_auth_config
|
||||
)
|
||||
response_part.function_response.id = call_part.function_call.id
|
||||
|
||||
final_response_parts = []
|
||||
async for event in runner.run_async(
|
||||
user_id="test_user",
|
||||
session_id=session.id,
|
||||
new_message=types.UserContent(parts=[response_part]),
|
||||
):
|
||||
event_list.append(event)
|
||||
if event.content:
|
||||
for part in event.content.parts:
|
||||
if part.text:
|
||||
final_response_parts.append(part.text)
|
||||
|
||||
final_response_text = "".join(final_response_parts)
|
||||
|
||||
# 6. Assertions
|
||||
|
||||
# Assert GCP Agent Identity client was invoked for credentials twice
|
||||
# (Initial Request + Post-Consent call)
|
||||
assert mock_client_cls.return_value.retrieve_credentials.call_count == 2
|
||||
expected_request = RetrieveCredentialsRequest(
|
||||
auth_provider=TEST_AUTH_PROVIDER_3LO,
|
||||
user_id="test_user",
|
||||
scopes=["test-scope"],
|
||||
continue_uri="https://example.com/continue",
|
||||
)
|
||||
mock_client_cls.return_value.retrieve_credentials.assert_called_with(
|
||||
expected_request
|
||||
)
|
||||
|
||||
assert "Tool executed successfully." in final_response_text
|
||||
|
||||
# Validate requests received by the mock model
|
||||
requests = mock_model.requests
|
||||
# Two LLM calls: the tool call in turn 1 (which ends at the EUC) and the
|
||||
# post-consent response in turn 2.
|
||||
assert len(requests) == 2
|
||||
|
||||
# Extract the function response from the prompt payload sent to the LLM
|
||||
last_request = requests[-1]
|
||||
function_response = next(
|
||||
(
|
||||
p.function_response
|
||||
for p in last_request.contents[-1].parts
|
||||
if p.function_response
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
assert function_response is not None
|
||||
assert function_response.name == "dummy_tool"
|
||||
assert DUMMY_TOKEN in str(function_response.response)
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright 2026 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.
|
||||
@@ -0,0 +1,57 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk.models.gemma_llm import Gemma
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.models.llm_response import LlmResponse
|
||||
from google.genai import types
|
||||
from google.genai.types import Content
|
||||
from google.genai.types import Part
|
||||
import pytest
|
||||
|
||||
DEFAULT_GEMMA_MODEL = "gemma-3-1b-it"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def gemma_llm():
|
||||
return Gemma(model=DEFAULT_GEMMA_MODEL)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def gemma_request():
|
||||
return LlmRequest(
|
||||
model=DEFAULT_GEMMA_MODEL,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part.from_text(text="You are a helpful assistant."),
|
||||
Part.from_text(text="Hello!"),
|
||||
],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
response_modalities=[types.Modality.TEXT],
|
||||
system_instruction="Talk like a pirate.",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"])
|
||||
async def test_generate_content_async(gemma_llm, gemma_request):
|
||||
async for response in gemma_llm.generate_content_async(gemma_request):
|
||||
assert isinstance(response, LlmResponse)
|
||||
assert response.content.parts[0].text
|
||||
@@ -0,0 +1,65 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk.models.google_llm import Gemini
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.models.llm_response import LlmResponse
|
||||
from google.genai import types
|
||||
from google.genai.types import Content
|
||||
from google.genai.types import Part
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def gemini_llm():
|
||||
return Gemini(model="gemini-2.5-flash")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llm_request():
|
||||
return LlmRequest(
|
||||
model="gemini-2.5-flash",
|
||||
contents=[Content(role="user", parts=[Part.from_text(text="Hello")])],
|
||||
config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
response_modalities=[types.Modality.TEXT],
|
||||
system_instruction="You are a helpful assistant",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async(gemini_llm, llm_request):
|
||||
async for response in gemini_llm.generate_content_async(llm_request):
|
||||
assert isinstance(response, LlmResponse)
|
||||
assert response.content.parts[0].text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async_stream(gemini_llm, llm_request):
|
||||
responses = [
|
||||
resp
|
||||
async for resp in gemini_llm.generate_content_async(
|
||||
llm_request, stream=True
|
||||
)
|
||||
]
|
||||
text = ""
|
||||
for i in range(len(responses) - 1):
|
||||
assert responses[i].partial is True
|
||||
assert responses[i].content.parts[0].text
|
||||
text += responses[i].content.parts[0].text
|
||||
|
||||
# Last message should be accumulated text
|
||||
assert responses[-1].content.parts[0].text == text
|
||||
assert not responses[-1].partial
|
||||
@@ -0,0 +1,165 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk.models.lite_llm import LiteLlm
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.models.llm_response import LlmResponse
|
||||
from google.genai import types
|
||||
from google.genai.types import Content
|
||||
from google.genai.types import Part
|
||||
import pytest
|
||||
|
||||
_TEST_MODEL_NAME = "vertex_ai/meta/llama-3.1-405b-instruct-maas"
|
||||
|
||||
_SYSTEM_PROMPT = """You are a helpful assistant."""
|
||||
|
||||
|
||||
def get_weather(city: str) -> str:
|
||||
"""Simulates a web search. Use it get information on weather.
|
||||
|
||||
Args:
|
||||
city: A string containing the location to get weather information for.
|
||||
|
||||
Returns:
|
||||
A string with the simulated weather information for the queried city.
|
||||
"""
|
||||
if "sf" in city.lower() or "san francisco" in city.lower():
|
||||
return "It's 70 degrees and foggy."
|
||||
return "It's 80 degrees and sunny."
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def oss_llm():
|
||||
return LiteLlm(model=_TEST_MODEL_NAME)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llm_request():
|
||||
return LlmRequest(
|
||||
model=_TEST_MODEL_NAME,
|
||||
contents=[Content(role="user", parts=[Part.from_text(text="hello")])],
|
||||
config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
response_modalities=[types.Modality.TEXT],
|
||||
system_instruction=_SYSTEM_PROMPT,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llm_request_with_tools():
|
||||
return LlmRequest(
|
||||
model=_TEST_MODEL_NAME,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part.from_text(text="What is the weather in San Francisco?")
|
||||
],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
response_modalities=[types.Modality.TEXT],
|
||||
system_instruction=_SYSTEM_PROMPT,
|
||||
tools=[
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name="get_weather",
|
||||
description="Get the weather in a given location",
|
||||
parameters=types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
properties={
|
||||
"city": types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description=(
|
||||
"The city to get the weather for."
|
||||
),
|
||||
),
|
||||
},
|
||||
required=["city"],
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async(oss_llm, llm_request):
|
||||
async for response in oss_llm.generate_content_async(llm_request):
|
||||
assert isinstance(response, LlmResponse)
|
||||
assert response.content.parts[0].text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async(oss_llm, llm_request):
|
||||
responses = [
|
||||
resp
|
||||
async for resp in oss_llm.generate_content_async(
|
||||
llm_request, stream=False
|
||||
)
|
||||
]
|
||||
part = responses[0].content.parts[0]
|
||||
assert len(part.text) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async_with_tools(
|
||||
oss_llm, llm_request_with_tools
|
||||
):
|
||||
responses = [
|
||||
resp
|
||||
async for resp in oss_llm.generate_content_async(
|
||||
llm_request_with_tools, stream=False
|
||||
)
|
||||
]
|
||||
function_call = responses[0].content.parts[0].function_call
|
||||
assert function_call.name == "get_weather"
|
||||
assert function_call.args["city"] == "San Francisco"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async_stream(oss_llm, llm_request):
|
||||
responses = [
|
||||
resp
|
||||
async for resp in oss_llm.generate_content_async(llm_request, stream=True)
|
||||
]
|
||||
text = ""
|
||||
for i in range(len(responses) - 1):
|
||||
assert responses[i].partial is True
|
||||
assert responses[i].content.parts[0].text
|
||||
text += responses[i].content.parts[0].text
|
||||
|
||||
# Last message should be accumulated text
|
||||
assert responses[-1].content.parts[0].text == text
|
||||
assert not responses[-1].partial
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async_stream_with_tools(
|
||||
oss_llm, llm_request_with_tools
|
||||
):
|
||||
responses = [
|
||||
resp
|
||||
async for resp in oss_llm.generate_content_async(
|
||||
llm_request_with_tools, stream=True
|
||||
)
|
||||
]
|
||||
function_call = responses[-1].content.parts[0].function_call
|
||||
assert function_call.name == "get_weather"
|
||||
assert function_call.args["city"] == "San Francisco"
|
||||
@@ -0,0 +1,112 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk.models.lite_llm import LiteLlm
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.genai import types
|
||||
from google.genai.types import Content
|
||||
from google.genai.types import Part
|
||||
import pytest
|
||||
|
||||
_TEST_MODEL_NAME = "vertex_ai/meta/llama-3.1-405b-instruct-maas"
|
||||
|
||||
_SYSTEM_PROMPT = """
|
||||
You are a helpful assistant, and call tools optionally.
|
||||
If call tools, the tool format should be in json body, and the tool argument values should be parsed from users inputs.
|
||||
"""
|
||||
|
||||
|
||||
_FUNCTIONS = [{
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The city to get the weather for.",
|
||||
},
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
}]
|
||||
|
||||
|
||||
def get_weather(city: str) -> str:
|
||||
"""Simulates a web search. Use it get information on weather.
|
||||
|
||||
Args:
|
||||
city: A string containing the location to get weather information for.
|
||||
|
||||
Returns:
|
||||
A string with the simulated weather information for the queried city.
|
||||
"""
|
||||
if "sf" in city.lower() or "san francisco" in city.lower():
|
||||
return "It's 70 degrees and foggy."
|
||||
return "It's 80 degrees and sunny."
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def oss_llm_with_function():
|
||||
return LiteLlm(model=_TEST_MODEL_NAME, functions=_FUNCTIONS)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llm_request():
|
||||
return LlmRequest(
|
||||
model=_TEST_MODEL_NAME,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part.from_text(text="What is the weather in San Francisco?")
|
||||
],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
response_modalities=[types.Modality.TEXT],
|
||||
system_instruction=_SYSTEM_PROMPT,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async_with_function(
|
||||
oss_llm_with_function, llm_request
|
||||
):
|
||||
responses = [
|
||||
resp
|
||||
async for resp in oss_llm_with_function.generate_content_async(
|
||||
llm_request, stream=False
|
||||
)
|
||||
]
|
||||
function_call = responses[0].content.parts[0].function_call
|
||||
assert function_call.name == "get_weather"
|
||||
assert function_call.args["city"] == "San Francisco"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async_stream_with_function(
|
||||
oss_llm_with_function, llm_request
|
||||
):
|
||||
responses = [
|
||||
resp
|
||||
async for resp in oss_llm_with_function.generate_content_async(
|
||||
llm_request, stream=True
|
||||
)
|
||||
]
|
||||
function_call = responses[-1].content.parts[0].function_call
|
||||
assert function_call.name == "get_weather"
|
||||
assert function_call.args["city"] == "San Francisco"
|
||||
@@ -0,0 +1,70 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from pytest import mark
|
||||
|
||||
from ..unittests.testing_utils import simplify_events
|
||||
from .fixture import callback_agent
|
||||
from .utils import assert_agent_says
|
||||
from .utils import TestRunner
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": callback_agent.agent.before_agent_callback_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_before_agent_call(agent_runner: TestRunner):
|
||||
agent_runner.run("Hi.")
|
||||
|
||||
# Assert the response content
|
||||
assert_agent_says(
|
||||
"End invocation event before agent call.",
|
||||
agent_name="before_agent_callback_agent",
|
||||
agent_runner=agent_runner,
|
||||
)
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": callback_agent.agent.before_model_callback_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_before_model_call(agent_runner: TestRunner):
|
||||
agent_runner.run("Hi.")
|
||||
|
||||
# Assert the response content
|
||||
assert_agent_says(
|
||||
"End invocation event before model call.",
|
||||
agent_name="before_model_callback_agent",
|
||||
agent_runner=agent_runner,
|
||||
)
|
||||
|
||||
|
||||
# TODO: re-enable vertex by removing below line after fixing.
|
||||
@mark.parametrize("llm_backend", ["GOOGLE_AI"], indirect=True)
|
||||
@mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": callback_agent.agent.after_model_callback_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_after_model_call(agent_runner: TestRunner):
|
||||
events = agent_runner.run("Hi.")
|
||||
|
||||
# Assert the response content
|
||||
simplified_events = simplify_events(events)
|
||||
assert simplified_events[0][0] == "after_model_callback_agent"
|
||||
assert simplified_events[0][1].endswith(
|
||||
"Update response event after model call."
|
||||
)
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copyright 2026 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 os
|
||||
from pathlib import Path
|
||||
|
||||
from click.testing import CliRunner
|
||||
from google.adk.cli import cli_tools_click
|
||||
import pytest
|
||||
|
||||
|
||||
def test_cli_run_integration(tmp_path: Path) -> None:
|
||||
"""Integration test for `adk run` with query using CliRunner (No mocks)."""
|
||||
# Arrange
|
||||
agent_dir = tmp_path / "dummy_agent"
|
||||
agent_dir.mkdir()
|
||||
|
||||
# Create __init__.py
|
||||
with open(agent_dir / "__init__.py", "w") as f:
|
||||
f.write("from . import agent\n")
|
||||
|
||||
# Create agent.py
|
||||
agent_code = """
|
||||
from google.adk.agents import Agent
|
||||
from google.adk.events.event import Event
|
||||
from typing import AsyncGenerator
|
||||
|
||||
class DummyAgent(Agent):
|
||||
async def run_async(self, ctx) -> AsyncGenerator[Event, None]:
|
||||
# Yield a text response
|
||||
from google.adk.events.event import Event
|
||||
text = ctx.user_content.parts[0].text if ctx.user_content and ctx.user_content.parts else "No input"
|
||||
yield Event(author="dummy", output=f"Echo: {text}")
|
||||
|
||||
root_agent = DummyAgent(name="dummy", model="gemini-2.5-flash")
|
||||
"""
|
||||
with open(agent_dir / "agent.py", "w") as f:
|
||||
f.write(agent_code)
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
# Act
|
||||
result = runner.invoke(
|
||||
cli_tools_click.main,
|
||||
["run", "--jsonl", str(agent_dir), "hello world"],
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result.exit_code == 0
|
||||
|
||||
# Check stdout
|
||||
stdout = result.stdout
|
||||
assert stdout, "Stdout should not be empty"
|
||||
|
||||
# Parse JSONL lines
|
||||
lines = stdout.strip().split("\n")
|
||||
assert len(lines) > 0
|
||||
|
||||
# The last line should be the final output or event
|
||||
last_line = lines[-1]
|
||||
try:
|
||||
data = json.loads(last_line)
|
||||
assert "output" in data
|
||||
assert "Echo: hello world" in data["output"]
|
||||
except json.JSONDecodeError:
|
||||
pytest.fail(f"Stdout contained non-JSON line: {last_line}")
|
||||
@@ -0,0 +1,67 @@
|
||||
# Copyright 2026 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 pytest
|
||||
|
||||
# Skip until fixed.
|
||||
pytest.skip(allow_module_level=True)
|
||||
|
||||
from .fixture import context_variable_agent
|
||||
from .utils import TestRunner
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": context_variable_agent.agent.state_variable_echo_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_context_variable_missing(agent_runner: TestRunner):
|
||||
with pytest.raises(KeyError) as e_info:
|
||||
agent_runner.run("Hi echo my customer id.")
|
||||
assert "customerId" in str(e_info.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": context_variable_agent.agent.state_variable_update_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_context_variable_update(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"update_fc",
|
||||
["RRRR", "3.141529", ["apple", "banana"], [1, 3.14, "hello"]],
|
||||
"successfully",
|
||||
)
|
||||
|
||||
|
||||
def _call_function_and_assert(
|
||||
agent_runner: TestRunner, function_name: str, params, expected
|
||||
):
|
||||
param_section = (
|
||||
" with params"
|
||||
f" {params if isinstance(params, str) else json.dumps(params)}"
|
||||
if params is not None
|
||||
else ""
|
||||
)
|
||||
agent_runner.run(
|
||||
f"Call {function_name}{param_section} and show me the result"
|
||||
)
|
||||
|
||||
model_response_event = agent_runner.get_events()[-1]
|
||||
assert model_response_event.author == "context_variable_update_agent"
|
||||
assert model_response_event.content.role == "model"
|
||||
assert expected in model_response_event.content.parts[0].text.strip()
|
||||
@@ -0,0 +1,71 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Evaluate all agents in fixture folder if evaluation test files exist."""
|
||||
|
||||
import os
|
||||
|
||||
from google.adk.evaluation.agent_evaluator import AgentEvaluator
|
||||
import pytest
|
||||
|
||||
|
||||
def agent_eval_artifacts_in_fixture():
|
||||
"""Get all agents from fixture folder."""
|
||||
agent_eval_artifacts = []
|
||||
fixture_dir = os.path.join(os.path.dirname(__file__), 'fixture')
|
||||
for agent_name in os.listdir(fixture_dir):
|
||||
agent_dir = os.path.join(fixture_dir, agent_name)
|
||||
if not os.path.isdir(agent_dir):
|
||||
continue
|
||||
for filename in os.listdir(agent_dir):
|
||||
# Evaluation test files end with test.json
|
||||
if not filename.endswith('test.json'):
|
||||
continue
|
||||
agent_eval_artifacts.append((
|
||||
f'tests.integration.fixture.{agent_name}',
|
||||
f'tests/integration/fixture/{agent_name}/{filename}',
|
||||
))
|
||||
|
||||
# This method gets invoked twice, sorting helps ensure that both the
|
||||
# invocations have the same view.
|
||||
agent_eval_artifacts = sorted(
|
||||
agent_eval_artifacts, key=lambda item: f'{item[0]}|{item[1]}'
|
||||
)
|
||||
return agent_eval_artifacts
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'agent_name, evalfile',
|
||||
agent_eval_artifacts_in_fixture(),
|
||||
ids=[agent_name for agent_name, _ in agent_eval_artifacts_in_fixture()],
|
||||
)
|
||||
async def test_evaluate_agents_long_running_4_runs_per_eval_item(
|
||||
agent_name, evalfile
|
||||
):
|
||||
"""Test agents evaluation in fixture folder.
|
||||
|
||||
After querying the fixture folder, we have 5 eval items. For each eval item
|
||||
we use 4 runs.
|
||||
|
||||
A single eval item is a session that can have multiple queries in it.
|
||||
"""
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module=agent_name,
|
||||
eval_dataset_file_path_or_dir=evalfile,
|
||||
# Using a slightly higher value helps us manage the variances that may
|
||||
# happen in each eval.
|
||||
# This, of course, comes at a cost of increased test run times.
|
||||
num_runs=4,
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Tests for generate_eval_cases CLI command."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
|
||||
from click.testing import CliRunner
|
||||
# We must mock or import the command safely for pytest
|
||||
from google.adk.cli.cli_tools_click import cli_generate_eval_cases
|
||||
import pytest
|
||||
|
||||
|
||||
def test_cli_generate_eval_cases_integration(tmp_path):
|
||||
"""E2E Test for the Vertex AI Scenario Generation Facade via the CLI."""
|
||||
# This requires identical project setup to Kokoro's e2e_test_gcp_ubuntu_docker
|
||||
if not os.environ.get("GOOGLE_CLOUD_PROJECT"):
|
||||
pytest.skip(
|
||||
"GOOGLE_CLOUD_PROJECT is not set. Skipping generation CLI integration"
|
||||
" test."
|
||||
)
|
||||
|
||||
# 1. Provide a UserSimulationConfig proxy
|
||||
config_file = tmp_path / "user_sim_config.json"
|
||||
config_data = {
|
||||
"generation_instruction": (
|
||||
"Generate a test conversation scenario where the user asks a simple"
|
||||
" question."
|
||||
),
|
||||
"count": 1,
|
||||
"model_name": "gemini-2.5-flash",
|
||||
}
|
||||
with open(config_file, "w") as f:
|
||||
json.dump(config_data, f)
|
||||
|
||||
eval_set_id = "cli_gen_test_eval_set"
|
||||
|
||||
# 2. Invoke the command via click's testing runner
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
cli_generate_eval_cases,
|
||||
[
|
||||
str(
|
||||
pathlib.Path(__file__).parent
|
||||
/ "fixture"
|
||||
/ "home_automation_agent"
|
||||
),
|
||||
eval_set_id,
|
||||
f"--user_simulation_config_file={config_file}",
|
||||
"--log_level=DEBUG",
|
||||
],
|
||||
)
|
||||
|
||||
# 3. Assert correct output
|
||||
assert (
|
||||
result.exit_code == 0
|
||||
), f"Command failed: {result.exception}\nOutput: {result.output}"
|
||||
assert "Generating scenarios utilizing Vertex AI Eval SDK..." in result.output
|
||||
assert "added to eval set" in result.output
|
||||
@@ -0,0 +1,172 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Live integration tests for ManagedAgent.
|
||||
|
||||
Assumes tests/integration/.env is present (auto-loaded by conftest.py) and that
|
||||
auth (ADC) is configured. Run explicitly:
|
||||
|
||||
pytest tests/integration/test_managed_agent.py -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from google.adk.agents import ManagedAgent
|
||||
from google.adk.runners import Runner
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.tools import google_search
|
||||
from google.adk.tools import RemoteMcpServer
|
||||
from google.adk.utils.context_utils import Aclosing
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
_AGENT_ID = 'antigravity-preview-05-2026'
|
||||
|
||||
|
||||
async def _run_turn(runner, session, text: str) -> list:
|
||||
events = []
|
||||
async with Aclosing(
|
||||
runner.run_async(
|
||||
user_id='test_user',
|
||||
session_id=session.id,
|
||||
new_message=types.Content(
|
||||
role='user', parts=[types.Part.from_text(text=text)]
|
||||
),
|
||||
)
|
||||
) as agen:
|
||||
async for event in agen:
|
||||
events.append(event)
|
||||
return events
|
||||
|
||||
|
||||
def _joined_text(events) -> str:
|
||||
return ' '.join(
|
||||
part.text
|
||||
for e in events
|
||||
if e.content and e.content.parts
|
||||
for part in e.content.parts
|
||||
if part.text
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_google_search_project_hail_mary():
|
||||
agent = ManagedAgent(
|
||||
name='managed_search_agent',
|
||||
agent_id=_AGENT_ID,
|
||||
environment={'type': 'remote'},
|
||||
tools=[google_search],
|
||||
)
|
||||
session_service = InMemorySessionService()
|
||||
runner = Runner(
|
||||
app_name='managed_agent_it',
|
||||
agent=agent,
|
||||
session_service=session_service,
|
||||
)
|
||||
session = await session_service.create_session(
|
||||
app_name='managed_agent_it', user_id='test_user'
|
||||
)
|
||||
|
||||
events = await _run_turn(
|
||||
runner, session, 'Who plays Rocky in the movie Project Hail Mary?'
|
||||
)
|
||||
|
||||
answer = _joined_text(events)
|
||||
print('\n=== ManagedAgent answer ===\n', answer)
|
||||
assert (
|
||||
'james ortiz' in answer.lower()
|
||||
), f'expected the grounded answer to contain "James Ortiz"; got: {answer!r}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_code_execution_prime_sum():
|
||||
agent = ManagedAgent(
|
||||
name='managed_code_execution_agent',
|
||||
agent_id=_AGENT_ID,
|
||||
environment={'type': 'remote'},
|
||||
tools=[types.Tool(code_execution=types.ToolCodeExecution())],
|
||||
)
|
||||
session_service = InMemorySessionService()
|
||||
runner = Runner(
|
||||
app_name='managed_agent_it',
|
||||
agent=agent,
|
||||
session_service=session_service,
|
||||
)
|
||||
session = await session_service.create_session(
|
||||
app_name='managed_agent_it', user_id='test_user'
|
||||
)
|
||||
|
||||
events = await _run_turn(
|
||||
runner,
|
||||
session,
|
||||
'What is the sum of the first 50 prime numbers? Use code to compute it.',
|
||||
)
|
||||
|
||||
answer = _joined_text(events)
|
||||
print('\n=== ManagedAgent code execution answer ===\n', answer)
|
||||
# The model may stream the number with thousands separators and/or stray
|
||||
# whitespace (e.g. "5,117" or "5, 117"), so remove all whitespace and commas
|
||||
# before matching the code-executed sum.
|
||||
normalized = ''.join(answer.split()).replace(',', '')
|
||||
assert (
|
||||
'5117' in normalized
|
||||
), f'expected the code-executed sum 5117; got: {answer!r}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
# Server-side remote MCP (mcp_server tool param) is currently only accepted by
|
||||
# the Gemini Developer API Interactions endpoint. The Vertex Interactions
|
||||
# endpoint returns 400 invalid_request for it (google-genai likewise documents
|
||||
# types.Tool.mcp_servers as unsupported on Vertex AI), so this live test is
|
||||
# scoped to GOOGLE_AI.
|
||||
@pytest.mark.parametrize('llm_backend', ['GOOGLE_AI'], indirect=True)
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get('GOOGLE_MAPS_API_KEY'),
|
||||
reason='GOOGLE_MAPS_API_KEY not set',
|
||||
)
|
||||
async def test_remote_mcp_maps_grounding_lite():
|
||||
agent = ManagedAgent(
|
||||
name='managed_maps_agent',
|
||||
agent_id=_AGENT_ID,
|
||||
environment={'type': 'remote'},
|
||||
tools=[
|
||||
RemoteMcpServer(
|
||||
name='maps_grounding_lite',
|
||||
url='https://mapstools.googleapis.com/mcp',
|
||||
header_provider=lambda ctx: {
|
||||
'X-Goog-Api-Key': os.environ['GOOGLE_MAPS_API_KEY']
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
session_service = InMemorySessionService()
|
||||
runner = Runner(
|
||||
app_name='managed_agent_it',
|
||||
agent=agent,
|
||||
session_service=session_service,
|
||||
)
|
||||
session = await session_service.create_session(
|
||||
app_name='managed_agent_it', user_id='test_user'
|
||||
)
|
||||
|
||||
events = await _run_turn(
|
||||
runner, session, 'Find a few coffee shops near Golden Gate Park.'
|
||||
)
|
||||
|
||||
# Non-deterministic content; assert a non-empty grounded answer came back and
|
||||
# no terminal error event was emitted.
|
||||
assert _joined_text(events).strip()
|
||||
assert not any(e.error_code for e in events)
|
||||
@@ -0,0 +1,27 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk.evaluation.agent_evaluator import AgentEvaluator
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_eval_agent():
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.trip_planner_agent",
|
||||
eval_dataset_file_path_or_dir=(
|
||||
"tests/integration/fixture/trip_planner_agent/trip_inquiry_multi_turn.test.json"
|
||||
),
|
||||
num_runs=4,
|
||||
)
|
||||
@@ -0,0 +1,46 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk.evaluation.agent_evaluator import AgentEvaluator
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simple_multi_turn_conversation():
|
||||
"""Test a simple multi-turn conversation."""
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.home_automation_agent",
|
||||
eval_dataset_file_path_or_dir="tests/integration/fixture/home_automation_agent/test_files/simple_multi_turn_conversation.test.json",
|
||||
num_runs=4,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dependent_tool_calls():
|
||||
"""Test subsequent tool calls that are dependent on previous tool calls."""
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.home_automation_agent",
|
||||
eval_dataset_file_path_or_dir="tests/integration/fixture/home_automation_agent/test_files/dependent_tool_calls.test.json",
|
||||
num_runs=4,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memorizing_past_events():
|
||||
"""Test memorizing past events."""
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.home_automation_agent",
|
||||
eval_dataset_file_path_or_dir="tests/integration/fixture/home_automation_agent/test_files/memorizing_past_events/eval_data.test.json",
|
||||
num_runs=4,
|
||||
)
|
||||
@@ -0,0 +1,45 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk.evaluation.agent_evaluator import AgentEvaluator
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_eval_agent():
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.home_automation_agent",
|
||||
eval_dataset_file_path_or_dir="tests/integration/fixture/home_automation_agent/simple_test.test.json",
|
||||
num_runs=4,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_eval_agent_with_agent_suffix_in_module_name():
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.home_automation_agent.agent",
|
||||
eval_dataset_file_path_or_dir="tests/integration/fixture/home_automation_agent/simple_test.test.json",
|
||||
num_runs=4,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_eval_agent_async():
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.hello_world_agent_async",
|
||||
eval_dataset_file_path_or_dir=(
|
||||
"tests/integration/fixture/hello_world_agent_async/roll_die.test.json"
|
||||
),
|
||||
num_runs=4,
|
||||
)
|
||||
@@ -0,0 +1,27 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk.evaluation.agent_evaluator import AgentEvaluator
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_eval_agent():
|
||||
"""Test hotel sub agent in a multi-agent system."""
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.trip_planner_agent",
|
||||
eval_dataset_file_path_or_dir="tests/integration/fixture/trip_planner_agent/test_files/trip_inquiry_sub_agent.test.json",
|
||||
agent_name="identify_agent",
|
||||
num_runs=4,
|
||||
)
|
||||
@@ -0,0 +1,177 @@
|
||||
# Copyright 2026 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 pytest
|
||||
|
||||
# Skip until fixed.
|
||||
pytest.skip(allow_module_level=True)
|
||||
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.sessions.session import Session
|
||||
from google.genai import types
|
||||
|
||||
from .fixture import context_variable_agent
|
||||
from .utils import TestRunner
|
||||
|
||||
nl_planner_si = """
|
||||
You are an intelligent tool use agent built upon the Gemini large language model. When answering the question, try to leverage the available tools to gather the information instead of your memorized knowledge.
|
||||
|
||||
Follow this process when answering the question: (1) first come up with a plan in natural language text format; (2) Then use tools to execute the plan and provide reasoning between tool code snippets to make a summary of current state and next step. Tool code snippets and reasoning should be interleaved with each other. (3) In the end, return one final answer.
|
||||
|
||||
Follow this format when answering the question: (1) The planning part should be under /*PLANNING*/. (2) The tool code snippets should be under /*ACTION*/, and the reasoning parts should be under /*REASONING*/. (3) The final answer part should be under /*FINAL_ANSWER*/.
|
||||
|
||||
|
||||
Below are the requirements for the planning:
|
||||
The plan is made to answer the user query if following the plan. The plan is coherent and covers all aspects of information from user query, and only involves the tools that are accessible by the agent. The plan contains the decomposed steps as a numbered list where each step should use one or multiple available tools. By reading the plan, you can intuitively know which tools to trigger or what actions to take.
|
||||
If the initial plan cannot be successfully executed, you should learn from previous execution results and revise your plan. The revised plan should be be under /*REPLANNING*/. Then use tools to follow the new plan.
|
||||
|
||||
Below are the requirements for the reasoning:
|
||||
The reasoning makes a summary of the current trajectory based on the user query and tool outputs. Based on the tool outputs and plan, the reasoning also comes up with instructions to the next steps, making the trajectory closer to the final answer.
|
||||
|
||||
|
||||
|
||||
Below are the requirements for the final answer:
|
||||
The final answer should be precise and follow query formatting requirements. Some queries may not be answerable with the available tools and information. In those cases, inform the user why you cannot process their query and ask for more information.
|
||||
|
||||
|
||||
|
||||
Below are the requirements for the tool code:
|
||||
|
||||
**Custom Tools:** The available tools are described in the context and can be directly used.
|
||||
- Code must be valid self-contained Python snippets with no imports and no references to tools or Python libraries that are not in the context.
|
||||
- You cannot use any parameters or fields that are not explicitly defined in the APIs in the context.
|
||||
- Use "print" to output execution results for the next step or final answer that you need for responding to the user. Never generate ```tool_outputs yourself.
|
||||
- The code snippets should be readable, efficient, and directly relevant to the user query and reasoning steps.
|
||||
- When using the tools, you should use the library name together with the function name, e.g., vertex_search.search().
|
||||
- If Python libraries are not provided in the context, NEVER write your own code other than the function calls using the provided tools.
|
||||
|
||||
|
||||
|
||||
VERY IMPORTANT instruction that you MUST follow in addition to the above instructions:
|
||||
|
||||
You should ask for clarification if you need more information to answer the question.
|
||||
You should prefer using the information available in the context instead of repeated tool use.
|
||||
|
||||
You should ONLY generate code snippets prefixed with "```tool_code" if you need to use the tools to answer the question.
|
||||
|
||||
If you are asked to write code by user specifically,
|
||||
- you should ALWAYS use "```python" to format the code.
|
||||
- you should NEVER put "tool_code" to format the code.
|
||||
- Good example:
|
||||
```python
|
||||
print('hello')
|
||||
```
|
||||
- Bad example:
|
||||
```tool_code
|
||||
print('hello')
|
||||
```
|
||||
"""
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": context_variable_agent.agent.state_variable_echo_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_context_variable(agent_runner: TestRunner):
|
||||
session = Session(
|
||||
context={
|
||||
"customerId": "1234567890",
|
||||
"customerInt": 30,
|
||||
"customerFloat": 12.34,
|
||||
"customerJson": {"name": "John Doe", "age": 30, "count": 11.1},
|
||||
}
|
||||
)
|
||||
si = UnitFlow()._build_system_instruction(
|
||||
InvocationContext(
|
||||
invocation_id="1234567890", agent=agent_runner.agent, session=session
|
||||
)
|
||||
)
|
||||
|
||||
assert (
|
||||
"Use the echo_info tool to echo 1234567890, 30, 12.34, and {'name': 'John"
|
||||
" Doe', 'age': 30, 'count': 11.1}. Ask for it if you need to."
|
||||
in si
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{
|
||||
"agent": (
|
||||
context_variable_agent.agent.state_variable_with_complicated_format_agent
|
||||
)
|
||||
}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_context_variable_with_complicated_format(agent_runner: TestRunner):
|
||||
session = Session(
|
||||
context={"customerId": "1234567890", "customer_int": 30},
|
||||
artifacts={"fileName": [types.Part(text="test artifact")]},
|
||||
)
|
||||
si = _context_formatter.populate_context_and_artifact_variable_values(
|
||||
agent_runner.agent.instruction,
|
||||
session.get_state(),
|
||||
session.get_artifact_dict(),
|
||||
)
|
||||
|
||||
assert (
|
||||
si
|
||||
== "Use the echo_info tool to echo 1234567890, 30, { "
|
||||
" non-identifier-float}}, test artifact, {'key1': 'value1'} and"
|
||||
" {{'key2': 'value2'}}. Ask for it if you need to."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{
|
||||
"agent": (
|
||||
context_variable_agent.agent.state_variable_with_nl_planner_agent
|
||||
)
|
||||
}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_nl_planner(agent_runner: TestRunner):
|
||||
session = Session(context={"customerId": "1234567890"})
|
||||
si = UnitFlow()._build_system_instruction(
|
||||
InvocationContext(
|
||||
invocation_id="1234567890",
|
||||
agent=agent_runner.agent,
|
||||
session=session,
|
||||
)
|
||||
)
|
||||
|
||||
for line in nl_planner_si.splitlines():
|
||||
assert line in si
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{
|
||||
"agent": (
|
||||
context_variable_agent.agent.state_variable_with_function_instruction_agent
|
||||
)
|
||||
}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_function_instruction(agent_runner: TestRunner):
|
||||
session = Session(context={"customerId": "1234567890"})
|
||||
si = UnitFlow()._build_system_instruction(
|
||||
InvocationContext(
|
||||
invocation_id="1234567890", agent=agent_runner.agent, session=session
|
||||
)
|
||||
)
|
||||
|
||||
assert "This is the plain text sub agent instruction." in si
|
||||
@@ -0,0 +1,287 @@
|
||||
# Copyright 2026 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 pytest
|
||||
|
||||
# Skip until fixed.
|
||||
pytest.skip(allow_module_level=True)
|
||||
|
||||
from .fixture import tool_agent
|
||||
from .utils import TestRunner
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.single_function_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_single_function_calls_success(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"simple_function",
|
||||
"test",
|
||||
"success",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_multiple_function_calls_success(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"simple_function",
|
||||
"test",
|
||||
"success",
|
||||
)
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"no_param_function",
|
||||
None,
|
||||
"Called no param function successfully",
|
||||
)
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"no_output_function",
|
||||
"test",
|
||||
"",
|
||||
)
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"multiple_param_types_function",
|
||||
["test", 1, 2.34, True],
|
||||
"success",
|
||||
)
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"return_list_str_function",
|
||||
"test",
|
||||
"success",
|
||||
)
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"list_str_param_function",
|
||||
["test", "test2", "test3", "test4"],
|
||||
"success",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Currently failing with 400 on MLDev.")
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_complex_function_calls_success(agent_runner: TestRunner):
|
||||
param1 = {"name": "Test", "count": 3}
|
||||
param2 = [
|
||||
{"name": "Function", "count": 2},
|
||||
{"name": "Retrieval", "count": 1},
|
||||
]
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"complex_function_list_dict",
|
||||
[param1, param2],
|
||||
"test",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_repetitive_call_success(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"repetitive_call_1",
|
||||
"test",
|
||||
"test_repetitive",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_function_calls_fail(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"throw_error_function",
|
||||
"test",
|
||||
None,
|
||||
ValueError,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_agent_tools_success(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"no_schema_agent",
|
||||
"Hi",
|
||||
"Hi",
|
||||
)
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"schema_agent",
|
||||
"Agent_tools",
|
||||
"Agent_tools_success",
|
||||
)
|
||||
_call_function_and_assert(
|
||||
agent_runner, "no_input_schema_agent", "Tools", "Tools_success"
|
||||
)
|
||||
_call_function_and_assert(agent_runner, "no_output_schema_agent", "Hi", "Hi")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_files_retrieval_success(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"test_case_retrieval",
|
||||
"What is the testing strategy of agent 2.0?",
|
||||
"test",
|
||||
)
|
||||
# For non relevant query, the agent should still be running fine, just return
|
||||
# response might be different for different calls, so we don't compare the
|
||||
# response here.
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"test_case_retrieval",
|
||||
"What is the whether in bay area?",
|
||||
"",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_rag_retrieval_success(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"valid_rag_retrieval",
|
||||
"What is the testing strategy of agent 2.0?",
|
||||
"test",
|
||||
)
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"valid_rag_retrieval",
|
||||
"What is the whether in bay area?",
|
||||
"No",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_rag_retrieval_fail(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"invalid_rag_retrieval",
|
||||
"What is the testing strategy of agent 2.0?",
|
||||
None,
|
||||
ValueError,
|
||||
)
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"non_exist_rag_retrieval",
|
||||
"What is the whether in bay area?",
|
||||
None,
|
||||
ValueError,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_langchain_tool_success(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"terminal",
|
||||
"Run the following shell command 'echo test!'",
|
||||
"test",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"agent_runner",
|
||||
[{"agent": tool_agent.agent.root_agent}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_crewai_tool_success(agent_runner: TestRunner):
|
||||
_call_function_and_assert(
|
||||
agent_runner,
|
||||
"directory_read_tool",
|
||||
"Find all the file paths",
|
||||
"file",
|
||||
)
|
||||
|
||||
|
||||
def _call_function_and_assert(
|
||||
agent_runner: TestRunner,
|
||||
function_name: str,
|
||||
params,
|
||||
expected=None,
|
||||
exception: Exception = None,
|
||||
):
|
||||
param_section = (
|
||||
" with params"
|
||||
f" {params if isinstance(params, str) else json.dumps(params)}"
|
||||
if params is not None
|
||||
else ""
|
||||
)
|
||||
query = f"Call {function_name}{param_section} and show me the result"
|
||||
if exception:
|
||||
_assert_raises(agent_runner, query, exception)
|
||||
return
|
||||
|
||||
_assert_function_output(agent_runner, query, expected)
|
||||
|
||||
|
||||
def _assert_raises(agent_runner: TestRunner, query: str, exception: Exception):
|
||||
with pytest.raises(exception):
|
||||
agent_runner.run(query)
|
||||
|
||||
|
||||
def _assert_function_output(agent_runner: TestRunner, query: str, expected):
|
||||
agent_runner.run(query)
|
||||
|
||||
# Retrieve the latest model response event
|
||||
model_response_event = agent_runner.get_events()[-1]
|
||||
|
||||
# Assert the response content
|
||||
assert model_response_event.content.role == "model"
|
||||
assert (
|
||||
expected.lower()
|
||||
in model_response_event.content.parts[0].text.strip().lower()
|
||||
)
|
||||
@@ -0,0 +1,334 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Integration tests for grounding metadata preservation in SSE streaming.
|
||||
|
||||
Verifies that grounding_metadata from VertexAiSearchTool reaches the final
|
||||
non-partial event in both progressive and non-progressive SSE streaming modes.
|
||||
|
||||
Prerequisites:
|
||||
- GOOGLE_CLOUD_PROJECT env var set to a GCP project with Vertex AI enabled
|
||||
- Discovery Engine API enabled (discoveryengine.googleapis.com)
|
||||
- Authenticated via `gcloud auth application-default login`
|
||||
|
||||
Usage:
|
||||
GOOGLE_CLOUD_PROJECT=my-project pytest
|
||||
tests/integration/test_vertex_ai_search_grounding_streaming.py -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from google.adk.features._feature_registry import FeatureName
|
||||
from google.adk.features._feature_registry import temporary_feature_override
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
_PROJECT = os.environ.get("GOOGLE_CLOUD_PROJECT", "")
|
||||
_LOCATION = os.environ.get("GOOGLE_CLOUD_LOCATION", "global")
|
||||
_COLLECTION = "default_collection"
|
||||
_DATA_STORE_ID = f"adk-grounding-test-{uuid.uuid4().hex[:8]}"
|
||||
_DATA_STORE_DISPLAY_NAME = "ADK Grounding Integration Test"
|
||||
_MODEL = "gemini-2.5-flash"
|
||||
|
||||
_TEST_DOCUMENTS = (
|
||||
{
|
||||
"id": "doc-adk-overview",
|
||||
"title": "ADK Overview",
|
||||
"content": (
|
||||
"The Agent Development Kit (ADK) is an open-source framework by"
|
||||
" Google for building AI agents. ADK supports multi-agent"
|
||||
" architectures, tool use, and integrates with Gemini models."
|
||||
" ADK was first released in April 2025."
|
||||
),
|
||||
},
|
||||
{
|
||||
"id": "doc-adk-tools",
|
||||
"title": "ADK Built-in Tools",
|
||||
"content": (
|
||||
"ADK provides built-in tools including VertexAiSearchTool for"
|
||||
" grounded search, GoogleSearchTool for web search, and"
|
||||
" CodeExecutionTool for running code. The VertexAiSearchTool"
|
||||
" returns grounding metadata with citations pointing to source"
|
||||
" documents."
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _parent_path() -> str:
|
||||
return f"projects/{_PROJECT}/locations/{_LOCATION}/collections/{_COLLECTION}"
|
||||
|
||||
|
||||
def _data_store_path() -> str:
|
||||
return f"{_parent_path()}/dataStores/{_DATA_STORE_ID}"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def project_id():
|
||||
if not _PROJECT:
|
||||
pytest.skip("GOOGLE_CLOUD_PROJECT env var not set")
|
||||
return _PROJECT
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def data_store_resource(project_id) -> str:
|
||||
"""Create a Vertex AI Search data store with test documents."""
|
||||
from google.api_core.exceptions import AlreadyExists
|
||||
from google.cloud import discoveryengine_v1beta as discoveryengine
|
||||
|
||||
ds_client = discoveryengine.DataStoreServiceClient()
|
||||
doc_client = discoveryengine.DocumentServiceClient()
|
||||
|
||||
# Create data store
|
||||
try:
|
||||
request = discoveryengine.CreateDataStoreRequest(
|
||||
parent=_parent_path(),
|
||||
data_store=discoveryengine.DataStore(
|
||||
display_name=_DATA_STORE_DISPLAY_NAME,
|
||||
industry_vertical=discoveryengine.IndustryVertical.GENERIC,
|
||||
solution_types=[discoveryengine.SolutionType.SOLUTION_TYPE_SEARCH],
|
||||
content_config=discoveryengine.DataStore.ContentConfig.NO_CONTENT,
|
||||
),
|
||||
data_store_id=_DATA_STORE_ID,
|
||||
)
|
||||
operation = ds_client.create_data_store(request=request)
|
||||
print(f"\nCreating data store '{_DATA_STORE_ID}'...")
|
||||
operation.result(timeout=120)
|
||||
print("Data store created.")
|
||||
except AlreadyExists:
|
||||
print(f"\nData store '{_DATA_STORE_ID}' already exists, reusing.")
|
||||
|
||||
# Ingest test documents
|
||||
branch = f"{_data_store_path()}/branches/default_branch"
|
||||
for doc_data in _TEST_DOCUMENTS:
|
||||
json_data = json.dumps({
|
||||
"title": doc_data["title"],
|
||||
"description": doc_data["content"],
|
||||
})
|
||||
doc = discoveryengine.Document(
|
||||
id=doc_data["id"],
|
||||
json_data=json_data,
|
||||
)
|
||||
try:
|
||||
doc_client.create_document(
|
||||
parent=branch,
|
||||
document=doc,
|
||||
document_id=doc_data["id"],
|
||||
)
|
||||
print(f" Created document: {doc_data['id']}")
|
||||
except AlreadyExists:
|
||||
doc_client.update_document(
|
||||
document=discoveryengine.Document(
|
||||
name=f"{branch}/documents/{doc_data['id']}",
|
||||
json_data=json_data,
|
||||
),
|
||||
)
|
||||
print(f" Updated document: {doc_data['id']}")
|
||||
|
||||
print("Waiting 5s for indexing...")
|
||||
time.sleep(5)
|
||||
|
||||
yield _data_store_path()
|
||||
|
||||
# Cleanup — best-effort, ignore errors from Discovery Engine LRO
|
||||
try:
|
||||
operation = ds_client.delete_data_store(name=_data_store_path())
|
||||
operation.result(timeout=120)
|
||||
print(f"\nDeleted data store '{_DATA_STORE_ID}'.")
|
||||
except Exception as e:
|
||||
print(f"\nFailed to delete data store '{_DATA_STORE_ID}': {e}")
|
||||
|
||||
|
||||
class TestIntegrationVertexAiSearchGrounding:
|
||||
"""Integration tests hitting real Vertex AI with VertexAiSearchTool."""
|
||||
|
||||
@pytest.mark.parametrize("llm_backend", ["VERTEX"], indirect=True)
|
||||
@pytest.mark.parametrize(
|
||||
"progressive_sse, label",
|
||||
[
|
||||
(True, "Progressive SSE"),
|
||||
(False, "Non-Progressive SSE"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_grounding_metadata_with_sse_streaming(
|
||||
self, project_id, data_store_resource, progressive_sse, label
|
||||
):
|
||||
"""Verifies grounding_metadata in SSE streaming modes."""
|
||||
from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.tools.vertex_ai_search_tool import VertexAiSearchTool
|
||||
|
||||
agent = LlmAgent(
|
||||
name="test_agent",
|
||||
model=_MODEL,
|
||||
tools=[VertexAiSearchTool(data_store_id=data_store_resource)],
|
||||
instruction="Answer questions using the search tool.",
|
||||
)
|
||||
|
||||
with temporary_feature_override(
|
||||
FeatureName.PROGRESSIVE_SSE_STREAMING, progressive_sse
|
||||
):
|
||||
all_events, saved_events = await self._run_agent_streaming(
|
||||
agent, project_id
|
||||
)
|
||||
|
||||
self._report_events(label, all_events, saved_events)
|
||||
|
||||
saved_with_grounding = [e for e in saved_events if e["has_grounding"]]
|
||||
assert (
|
||||
saved_with_grounding
|
||||
), f"No saved (non-partial) events have grounding_metadata with {label}."
|
||||
|
||||
@pytest.mark.parametrize("llm_backend", ["VERTEX"], indirect=True)
|
||||
@pytest.mark.asyncio
|
||||
async def test_grounding_metadata_without_streaming(
|
||||
self, project_id, data_store_resource
|
||||
):
|
||||
"""Without streaming, grounding_metadata should always be present."""
|
||||
from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.agents.run_config import RunConfig
|
||||
from google.adk.agents.run_config import StreamingMode
|
||||
from google.adk.runners import Runner
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.tools.vertex_ai_search_tool import VertexAiSearchTool
|
||||
from google.adk.utils.context_utils import Aclosing
|
||||
|
||||
agent = LlmAgent(
|
||||
name="test_agent",
|
||||
model=_MODEL,
|
||||
tools=[VertexAiSearchTool(data_store_id=data_store_resource)],
|
||||
instruction="Answer questions using the search tool.",
|
||||
)
|
||||
|
||||
session_service = InMemorySessionService()
|
||||
runner = Runner(
|
||||
app_name="test_app",
|
||||
agent=agent,
|
||||
session_service=session_service,
|
||||
)
|
||||
session = await session_service.create_session(
|
||||
app_name="test_app", user_id="test_user"
|
||||
)
|
||||
|
||||
run_config = RunConfig(streaming_mode=StreamingMode.NONE)
|
||||
events = []
|
||||
async with Aclosing(
|
||||
runner.run_async(
|
||||
user_id="test_user",
|
||||
session_id=session.id,
|
||||
new_message=types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
types.Part.from_text(
|
||||
text="What built-in tools does ADK provide?"
|
||||
)
|
||||
],
|
||||
),
|
||||
run_config=run_config,
|
||||
)
|
||||
) as agen:
|
||||
async for event in agen:
|
||||
events.append({
|
||||
"author": event.author,
|
||||
"partial": event.partial,
|
||||
"has_grounding": event.grounding_metadata is not None,
|
||||
"has_content": bool(event.content and event.content.parts),
|
||||
})
|
||||
|
||||
print("\n=== No Streaming ===")
|
||||
for i, e in enumerate(events):
|
||||
print(
|
||||
f" Event {i}: author={e['author']}, partial={e['partial']},"
|
||||
f" grounding={e['has_grounding']}, content={e['has_content']}"
|
||||
)
|
||||
|
||||
model_events = [e for e in events if e["author"] == "test_agent"]
|
||||
with_grounding = [e for e in model_events if e["has_grounding"]]
|
||||
assert (
|
||||
with_grounding
|
||||
), "No events have grounding_metadata even without streaming."
|
||||
|
||||
async def _run_agent_streaming(self, agent, project_id):
|
||||
from google.adk.agents.run_config import RunConfig
|
||||
from google.adk.agents.run_config import StreamingMode
|
||||
from google.adk.runners import Runner
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.utils.context_utils import Aclosing
|
||||
|
||||
session_service = InMemorySessionService()
|
||||
runner = Runner(
|
||||
app_name="test_app",
|
||||
agent=agent,
|
||||
session_service=session_service,
|
||||
)
|
||||
session = await session_service.create_session(
|
||||
app_name="test_app", user_id="test_user"
|
||||
)
|
||||
|
||||
run_config = RunConfig(streaming_mode=StreamingMode.SSE)
|
||||
all_events = []
|
||||
async with Aclosing(
|
||||
runner.run_async(
|
||||
user_id="test_user",
|
||||
session_id=session.id,
|
||||
new_message=types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
types.Part.from_text(
|
||||
text="What is ADK and when was it first released?"
|
||||
)
|
||||
],
|
||||
),
|
||||
run_config=run_config,
|
||||
)
|
||||
) as agen:
|
||||
async for event in agen:
|
||||
all_events.append({
|
||||
"author": event.author,
|
||||
"partial": event.partial,
|
||||
"has_grounding": event.grounding_metadata is not None,
|
||||
"has_content": bool(event.content and event.content.parts),
|
||||
})
|
||||
|
||||
saved_events = [e for e in all_events if e["partial"] is not True]
|
||||
return all_events, saved_events
|
||||
|
||||
def _report_events(self, label, all_events, saved_events):
|
||||
print(f"\n=== {label} — All Events ===")
|
||||
for i, e in enumerate(all_events):
|
||||
print(
|
||||
f" Event {i}: author={e['author']}, partial={e['partial']},"
|
||||
f" grounding={e['has_grounding']},"
|
||||
f" content={e['has_content']}"
|
||||
)
|
||||
print(f"\n=== {label} — Saved (non-partial) Events ===")
|
||||
for i, e in enumerate(saved_events):
|
||||
print(
|
||||
f" Event {i}: author={e['author']}, partial={e['partial']},"
|
||||
f" grounding={e['has_grounding']},"
|
||||
f" content={e['has_content']}"
|
||||
)
|
||||
partial_with_grounding = [
|
||||
e for e in all_events if e["partial"] is True and e["has_grounding"]
|
||||
]
|
||||
if partial_with_grounding:
|
||||
print(
|
||||
f"\n NOTE: {len(partial_with_grounding)} partial event(s)"
|
||||
" had grounding_metadata but were NOT saved to session."
|
||||
)
|
||||
@@ -0,0 +1,37 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from google.adk.evaluation.agent_evaluator import AgentEvaluator
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_single_test_file():
|
||||
"""Test the agent's basic ability via session file."""
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.home_automation_agent",
|
||||
eval_dataset_file_path_or_dir="tests/integration/fixture/home_automation_agent/simple_test.test.json",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_folder_of_test_files_long_running():
|
||||
"""Test the agent's basic ability via a folder of session files."""
|
||||
await AgentEvaluator.evaluate(
|
||||
agent_module="tests.integration.fixture.home_automation_agent",
|
||||
eval_dataset_file_path_or_dir=(
|
||||
"tests/integration/fixture/home_automation_agent/test_files"
|
||||
),
|
||||
num_runs=4,
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright 2026 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.
|
||||
@@ -0,0 +1,16 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from .asserts import *
|
||||
from .test_runner import TestRunner
|
||||
@@ -0,0 +1,75 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
from typing import TypedDict
|
||||
|
||||
from .test_runner import TestRunner
|
||||
|
||||
|
||||
class Message(TypedDict):
|
||||
agent_name: str
|
||||
expected_text: str
|
||||
|
||||
|
||||
def assert_current_agent_is(agent_name: str, *, agent_runner: TestRunner):
|
||||
assert agent_runner.get_current_agent_name() == agent_name
|
||||
|
||||
|
||||
def assert_agent_says(
|
||||
expected_text: str, *, agent_name: str, agent_runner: TestRunner
|
||||
):
|
||||
for event in reversed(agent_runner.get_events()):
|
||||
if event.author == agent_name and event.content.parts[0].text:
|
||||
assert event.content.parts[0].text.strip() == expected_text
|
||||
return
|
||||
|
||||
|
||||
def assert_agent_says_in_order(
|
||||
expected_conversation: list[Message], agent_runner: TestRunner
|
||||
):
|
||||
expected_conversation_idx = len(expected_conversation) - 1
|
||||
for event in reversed(agent_runner.get_events()):
|
||||
if event.content.parts and event.content.parts[0].text:
|
||||
assert (
|
||||
event.author
|
||||
== expected_conversation[expected_conversation_idx]['agent_name']
|
||||
)
|
||||
assert (
|
||||
event.content.parts[0].text.strip()
|
||||
== expected_conversation[expected_conversation_idx]['expected_text']
|
||||
)
|
||||
expected_conversation_idx -= 1
|
||||
if expected_conversation_idx < 0:
|
||||
return
|
||||
|
||||
|
||||
def assert_agent_transfer_path(
|
||||
expected_path: list[str], *, agent_runner: TestRunner
|
||||
):
|
||||
events = agent_runner.get_events()
|
||||
idx_in_expected_path = len(expected_path) - 1
|
||||
# iterate events in reverse order
|
||||
for event in reversed(events):
|
||||
function_calls = event.get_function_calls()
|
||||
if (
|
||||
len(function_calls) == 1
|
||||
and function_calls[0].name == 'transfer_to_agent'
|
||||
):
|
||||
assert (
|
||||
function_calls[0].args['agent_name']
|
||||
== expected_path[idx_in_expected_path]
|
||||
)
|
||||
idx_in_expected_path -= 1
|
||||
if idx_in_expected_path < 0:
|
||||
return
|
||||
@@ -0,0 +1,97 @@
|
||||
# Copyright 2026 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 importlib
|
||||
from typing import Optional
|
||||
|
||||
from google.adk import Agent
|
||||
from google.adk import Runner
|
||||
from google.adk.artifacts.base_artifact_service import BaseArtifactService
|
||||
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||||
from google.adk.events.event import Event
|
||||
from google.adk.sessions.base_session_service import BaseSessionService
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.sessions.session import Session
|
||||
from google.genai import types
|
||||
|
||||
|
||||
class TestRunner:
|
||||
"""Agents runner for testing."""
|
||||
|
||||
app_name = "test_app"
|
||||
user_id = "test_user"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: Agent,
|
||||
artifact_service: BaseArtifactService = InMemoryArtifactService(),
|
||||
session_service: BaseSessionService = InMemorySessionService(),
|
||||
) -> None:
|
||||
self.agent = agent
|
||||
self.agent_client = Runner(
|
||||
app_name=self.app_name,
|
||||
agent=agent,
|
||||
artifact_service=artifact_service,
|
||||
session_service=session_service,
|
||||
)
|
||||
self.session_service = session_service
|
||||
self.current_session_id = session_service.create_session(
|
||||
app_name=self.app_name, user_id=self.user_id
|
||||
).id
|
||||
|
||||
def new_session(self, session_id: Optional[str] = None) -> None:
|
||||
self.current_session_id = self.session_service.create_session(
|
||||
app_name=self.app_name, user_id=self.user_id, session_id=session_id
|
||||
).id
|
||||
|
||||
def run(self, prompt: str) -> list[Event]:
|
||||
current_session = self.session_service.get_session(
|
||||
app_name=self.app_name,
|
||||
user_id=self.user_id,
|
||||
session_id=self.current_session_id,
|
||||
)
|
||||
assert current_session is not None
|
||||
|
||||
return list(
|
||||
self.agent_client.run(
|
||||
user_id=current_session.user_id,
|
||||
session_id=current_session.id,
|
||||
new_message=types.Content(
|
||||
role="user",
|
||||
parts=[types.Part.from_text(text=prompt)],
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
def get_current_session(self) -> Optional[Session]:
|
||||
return self.session_service.get_session(
|
||||
app_name=self.app_name,
|
||||
user_id=self.user_id,
|
||||
session_id=self.current_session_id,
|
||||
)
|
||||
|
||||
def get_events(self) -> list[Event]:
|
||||
return self.get_current_session().events
|
||||
|
||||
@classmethod
|
||||
def from_agent_name(cls, agent_name: str):
|
||||
agent_module_path = f"tests.integration.fixture.{agent_name}"
|
||||
agent_module = importlib.import_module(agent_module_path)
|
||||
agent: Agent = agent_module.agent.root_agent
|
||||
return cls(agent)
|
||||
|
||||
def get_current_agent_name(self) -> str:
|
||||
return self.agent_client._find_agent_to_run(
|
||||
self.get_current_session(), self.agent
|
||||
).name
|
||||
@@ -0,0 +1,67 @@
|
||||
# Remote Integration Tests for `/apps/{app_name}/trigger/*` Endpoints
|
||||
|
||||
End-to-end tests that verify Pub/Sub push, and
|
||||
Eventarc triggers against a real ADK agent deployed to Cloud Run.
|
||||
|
||||
## Workflow
|
||||
|
||||
This workflow is split into three independent phases.
|
||||
|
||||
### Phase 1: Deploy Agent
|
||||
|
||||
Build the ADK package and deploy the echo agent to Cloud Run. This ensures that
|
||||
the service and its identity exist before any infrastructure wiring occurs.
|
||||
|
||||
```bash
|
||||
export GCP_PROJECT_ID=your-project-id
|
||||
export SUFFIX=ea786 # Pick a unique suffix
|
||||
export SERVICE_NAME=adk-trigger-test-$SUFFIX
|
||||
|
||||
# 1. Build local ADK wheel
|
||||
uv build --wheel --out-dir tests/remote/triggers/test_agent/wheels/
|
||||
|
||||
# 2. Deploy Agent
|
||||
gcloud run deploy $SERVICE_NAME \
|
||||
--source=tests/remote/triggers/test_agent \
|
||||
--project="$GCP_PROJECT_ID" \
|
||||
--region="us-central1" \
|
||||
--port=8080 \
|
||||
--quiet
|
||||
```
|
||||
|
||||
### Phase 2: Wire Infrastructure (Terraform)
|
||||
|
||||
Run Terraform to create the supporting infrastructure (IAM roles, Pub/Sub
|
||||
topics).
|
||||
|
||||
```bash
|
||||
cd tests/remote/triggers/terraform
|
||||
terraform init
|
||||
terraform apply \
|
||||
-var=project_id=$GCP_PROJECT_ID \
|
||||
-var=service_name=$SERVICE_NAME \
|
||||
-var=suffix=$SUFFIX
|
||||
```
|
||||
|
||||
### Phase 3: Run Tests (Pytest)
|
||||
|
||||
```bash
|
||||
# Run Tests from the project root
|
||||
export GCP_PROJECT_ID=your-project-id
|
||||
export SUFFIX=ea786
|
||||
pytest tests/remote/triggers/ -v -s
|
||||
```
|
||||
|
||||
## Cleanup
|
||||
|
||||
```bash
|
||||
# 1. Destroy infrastructure
|
||||
cd tests/remote/triggers/terraform
|
||||
terraform destroy \
|
||||
-var=project_id=$GCP_PROJECT_ID \
|
||||
-var=service_name=$SERVICE_NAME \
|
||||
-var=suffix=$SUFFIX
|
||||
|
||||
# 2. Delete Cloud Run service
|
||||
gcloud run services delete $SERVICE_NAME --project=$GCP_PROJECT_ID --region=us-central1 --quiet
|
||||
```
|
||||
@@ -0,0 +1,169 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Shared pytest fixtures for remote trigger integration tests.
|
||||
|
||||
Exposes GCP resource references by reading current Terraform state.
|
||||
|
||||
Environment variables:
|
||||
GCP_PROJECT_ID : GCP project.
|
||||
GCP_REGION : GCP region (default: ``us-central1``).
|
||||
ADK_TERRAFORM_BIN : Path to terraform binary (default: ``terraform``).
|
||||
ADK_TERRAFORM_CWD : Directory to run terraform from (default:
|
||||
``tests/remote/terraform``).
|
||||
ADK_TERRAFORM_ARGS: Extra arguments for terraform commands.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
TERRAFORM_DIR = os.path.join(os.path.dirname(__file__), "terraform")
|
||||
|
||||
|
||||
def _get_project_id() -> str | None:
|
||||
"""Return GCP project ID from env or gcloud config."""
|
||||
project = os.environ.get("GCP_PROJECT_ID")
|
||||
if project:
|
||||
return project
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["gcloud", "config", "get-value", "project"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
return result.stdout.strip() or None
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
|
||||
|
||||
def _get_identity_token(audience: str) -> str | None:
|
||||
"""Fetch an identity token for the given audience via gcloud."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"gcloud",
|
||||
"auth",
|
||||
"print-identity-token",
|
||||
f"--audiences={audience}",
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return result.stdout.strip()
|
||||
except (FileNotFoundError, subprocess.CalledProcessError):
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Infrastructure Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def terraform_outputs():
|
||||
"""Read Terraform outputs from the current state."""
|
||||
project = _get_project_id()
|
||||
if not project:
|
||||
pytest.skip(
|
||||
"GCP_PROJECT_ID not set and no gcloud default project configured"
|
||||
)
|
||||
|
||||
tf_bin = os.environ.get("ADK_TERRAFORM_BIN", "terraform")
|
||||
tf_cwd = os.environ.get("ADK_TERRAFORM_CWD", TERRAFORM_DIR)
|
||||
tf_args = shlex.split(os.environ.get("ADK_TERRAFORM_ARGS", ""))
|
||||
|
||||
try:
|
||||
# Build the command using provided overrides
|
||||
cmd = [tf_bin] + tf_args + ["output", "-json"]
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
cwd=tf_cwd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
raw = json.loads(result.stdout)
|
||||
return {k: v["value"] for k, v in raw.items()}
|
||||
except (
|
||||
subprocess.CalledProcessError,
|
||||
FileNotFoundError,
|
||||
json.JSONDecodeError,
|
||||
) as e:
|
||||
pytest.fail(
|
||||
"Failed to read Terraform outputs. Ensure 'terraform apply' has been"
|
||||
f" run successfully.\nCommand: {' '.join(cmd)}\nCWD:"
|
||||
f" {tf_cwd}\nError: {e}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def cloud_run_url(terraform_outputs) -> str:
|
||||
"""Base URL of the deployed Cloud Run service."""
|
||||
return terraform_outputs["cloud_run_url"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def pubsub_topic(terraform_outputs) -> str:
|
||||
"""Fully qualified Pub/Sub topic name."""
|
||||
return terraform_outputs["pubsub_topic"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def pubsub_topic_short(terraform_outputs) -> str:
|
||||
"""Short Pub/Sub topic name."""
|
||||
return terraform_outputs["pubsub_topic_short"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def eventarc_topic(terraform_outputs) -> str:
|
||||
"""Fully qualified Eventarc source topic name."""
|
||||
return terraform_outputs["eventarc_topic"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def eventarc_topic_short(terraform_outputs) -> str:
|
||||
"""Short Eventarc source topic name."""
|
||||
return terraform_outputs["eventarc_topic_short"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def project_id(terraform_outputs) -> str:
|
||||
"""GCP project ID."""
|
||||
return terraform_outputs["project_id"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def auth_headers(cloud_run_url) -> dict[str, str]:
|
||||
"""Authorization headers for direct HTTP calls to Cloud Run."""
|
||||
try:
|
||||
resp = requests.get(cloud_run_url, timeout=10)
|
||||
if resp.status_code != 403:
|
||||
return {}
|
||||
except requests.RequestException:
|
||||
pass
|
||||
|
||||
token = _get_identity_token(cloud_run_url)
|
||||
if token:
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
return {}
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ADK Agent for trigger testing.
|
||||
#
|
||||
# This configuration references a Cloud Run service that has been deployed
|
||||
# manually (e.g. via `gcloud run deploy`) before running Terraform.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
locals {
|
||||
service_name = var.service_name != null ? var.service_name : "${local.name_prefix}-${local.suffix}"
|
||||
}
|
||||
|
||||
# Read the service back as a data source so other resources can reference
|
||||
# its URL and attributes.
|
||||
data "google_cloud_run_v2_service" "trigger_agent" {
|
||||
name = local.service_name
|
||||
location = var.region
|
||||
project = var.project_id
|
||||
}
|
||||
|
||||
# No longer using resource "google_cloud_run_v2_service" to avoid
|
||||
# deployment conflicts with local tools.
|
||||
@@ -0,0 +1,120 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Eventarc: separate Pub/Sub topic as event source → Cloud Run /trigger/eventarc
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# A dedicated topic that acts as the Eventarc event source.
|
||||
# Publishing to this topic triggers the Eventarc → Cloud Run pipeline.
|
||||
resource "google_pubsub_topic" "eventarc_source" {
|
||||
name = "${local.name_prefix}-eventarc-${local.suffix}"
|
||||
project = var.project_id
|
||||
|
||||
depends_on = [google_project_service.apis]
|
||||
}
|
||||
|
||||
resource "google_eventarc_trigger" "trigger_test" {
|
||||
name = "${local.name_prefix}-${local.suffix}"
|
||||
location = var.region
|
||||
project = var.project_id
|
||||
|
||||
matching_criteria {
|
||||
attribute = "type"
|
||||
value = "google.cloud.pubsub.topic.v1.messagePublished"
|
||||
}
|
||||
|
||||
transport {
|
||||
pubsub {
|
||||
topic = google_pubsub_topic.eventarc_source.id
|
||||
}
|
||||
}
|
||||
|
||||
destination {
|
||||
cloud_run_service {
|
||||
service = data.google_cloud_run_v2_service.trigger_agent.name
|
||||
path = "/apps/trigger_echo_agent/trigger/eventarc"
|
||||
region = var.region
|
||||
}
|
||||
}
|
||||
|
||||
service_account = google_service_account.eventarc_invoker.email
|
||||
|
||||
depends_on = [
|
||||
google_project_iam_member.eventarc_event_receiver,
|
||||
google_project_service.apis,
|
||||
]
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GCS Trigger Test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
resource "random_id" "bucket_suffix" {
|
||||
byte_length = 4
|
||||
}
|
||||
|
||||
resource "google_storage_bucket" "trigger_test_bucket" {
|
||||
name = "${local.name_prefix}-bucket-${local.suffix}-${random_id.bucket_suffix.hex}"
|
||||
location = var.region
|
||||
project = var.project_id
|
||||
force_destroy = true
|
||||
|
||||
uniform_bucket_level_access = true
|
||||
|
||||
depends_on = [google_project_service.apis]
|
||||
}
|
||||
|
||||
data "google_storage_project_service_account" "gcs_account" {
|
||||
project = var.project_id
|
||||
}
|
||||
|
||||
resource "google_project_iam_member" "gcs_pubsub_publisher" {
|
||||
project = var.project_id
|
||||
role = "roles/pubsub.publisher"
|
||||
member = "serviceAccount:${data.google_storage_project_service_account.gcs_account.email_address}"
|
||||
|
||||
depends_on = [data.google_storage_project_service_account.gcs_account]
|
||||
}
|
||||
|
||||
resource "google_eventarc_trigger" "gcs_trigger" {
|
||||
name = "${local.name_prefix}-gcs-${local.suffix}"
|
||||
location = var.region
|
||||
project = var.project_id
|
||||
|
||||
matching_criteria {
|
||||
attribute = "type"
|
||||
value = "google.cloud.storage.object.v1.finalized"
|
||||
}
|
||||
matching_criteria {
|
||||
attribute = "bucket"
|
||||
value = google_storage_bucket.trigger_test_bucket.name
|
||||
}
|
||||
|
||||
destination {
|
||||
cloud_run_service {
|
||||
service = data.google_cloud_run_v2_service.trigger_agent.name
|
||||
path = "/apps/trigger_echo_agent/trigger/eventarc"
|
||||
region = var.region
|
||||
}
|
||||
}
|
||||
|
||||
service_account = google_service_account.eventarc_invoker.email
|
||||
|
||||
depends_on = [
|
||||
google_project_iam_member.eventarc_event_receiver,
|
||||
google_project_iam_member.gcs_pubsub_publisher,
|
||||
google_project_service.apis,
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# IAM bindings and Service Accounts for Cloud Run invokers.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Service account for the Cloud Run agent itself.
|
||||
resource "google_service_account" "cloud_run" {
|
||||
account_id = "${local.name_prefix}-run-${local.suffix}"
|
||||
display_name = "ADK Trigger Test - Cloud Run Agent"
|
||||
project = var.project_id
|
||||
}
|
||||
|
||||
# Service account for Pub/Sub to invoke Cloud Run.
|
||||
resource "google_service_account" "pubsub_invoker" {
|
||||
account_id = "${local.name_prefix}-ps-${local.suffix}"
|
||||
display_name = "ADK Trigger Test - Pub/Sub Invoker"
|
||||
project = var.project_id
|
||||
}
|
||||
|
||||
# Service account for Eventarc to invoke Cloud Run.
|
||||
resource "google_service_account" "eventarc_invoker" {
|
||||
account_id = "${local.name_prefix}-ea-${local.suffix}"
|
||||
display_name = "ADK Trigger Test - Eventarc Invoker"
|
||||
project = var.project_id
|
||||
}
|
||||
|
||||
resource "google_cloud_run_v2_service_iam_member" "pubsub_invoker" {
|
||||
name = data.google_cloud_run_v2_service.trigger_agent.name
|
||||
location = var.region
|
||||
project = var.project_id
|
||||
role = "roles/run.invoker"
|
||||
member = "serviceAccount:${google_service_account.pubsub_invoker.email}"
|
||||
}
|
||||
|
||||
resource "google_cloud_run_v2_service_iam_member" "eventarc_invoker" {
|
||||
name = data.google_cloud_run_v2_service.trigger_agent.name
|
||||
location = var.region
|
||||
project = var.project_id
|
||||
role = "roles/run.invoker"
|
||||
member = "serviceAccount:${google_service_account.eventarc_invoker.email}"
|
||||
}
|
||||
|
||||
|
||||
# Eventarc requires the receiver role on the invoker service account.
|
||||
resource "google_project_iam_member" "eventarc_event_receiver" {
|
||||
project = var.project_id
|
||||
role = "roles/eventarc.eventReceiver"
|
||||
member = "serviceAccount:${google_service_account.eventarc_invoker.email}"
|
||||
}
|
||||
|
||||
data "google_project" "project" {
|
||||
project_id = var.project_id
|
||||
}
|
||||
|
||||
# Grant the Pub/Sub service agent permission to create OIDC tokens for the pubsub_invoker SA
|
||||
resource "google_service_account_iam_member" "pubsub_token_creator" {
|
||||
service_account_id = google_service_account.pubsub_invoker.name
|
||||
role = "roles/iam.serviceAccountTokenCreator"
|
||||
member = "serviceAccount:service-${data.google_project.project.number}@gcp-sa-pubsub.iam.gserviceaccount.com"
|
||||
}
|
||||
|
||||
# Grant the Pub/Sub service agent permission to create OIDC tokens for the eventarc_invoker SA
|
||||
resource "google_service_account_iam_member" "eventarc_token_creator" {
|
||||
service_account_id = google_service_account.eventarc_invoker.name
|
||||
role = "roles/iam.serviceAccountTokenCreator"
|
||||
member = "serviceAccount:service-${data.google_project.project.number}@gcp-sa-pubsub.iam.gserviceaccount.com"
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
terraform {
|
||||
required_version = ">= 1.0.0"
|
||||
|
||||
required_providers {
|
||||
google = {
|
||||
source = "hashicorp/google"
|
||||
version = ">= 5.0.0"
|
||||
}
|
||||
random = {
|
||||
source = "hashicorp/random"
|
||||
version = ">= 3.0.0"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
provider "google" {
|
||||
project = var.project_id
|
||||
region = var.region
|
||||
|
||||
default_labels = {
|
||||
goog-terraform-provisioned = "true"
|
||||
}
|
||||
}
|
||||
|
||||
# Random suffix to avoid resource name collisions across parallel test runs.
|
||||
resource "random_id" "suffix" {
|
||||
byte_length = 4
|
||||
}
|
||||
|
||||
locals {
|
||||
suffix = var.suffix != null ? var.suffix : random_id.suffix.hex
|
||||
name_prefix = "adk-trigger-test"
|
||||
|
||||
required_apis = [
|
||||
"run.googleapis.com",
|
||||
"pubsub.googleapis.com",
|
||||
"cloudbuild.googleapis.com",
|
||||
"eventarc.googleapis.com",
|
||||
"logging.googleapis.com",
|
||||
]
|
||||
}
|
||||
|
||||
resource "google_project_service" "apis" {
|
||||
for_each = toset(local.required_apis)
|
||||
|
||||
project = var.project_id
|
||||
service = each.value
|
||||
|
||||
disable_on_destroy = false
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
output "cloud_run_url" {
|
||||
description = "Base URL of the deployed Cloud Run trigger-echo-agent service."
|
||||
value = data.google_cloud_run_v2_service.trigger_agent.uri
|
||||
}
|
||||
|
||||
output "pubsub_topic" {
|
||||
description = "Fully qualified Pub/Sub topic name for direct-push tests."
|
||||
value = google_pubsub_topic.trigger_test.id
|
||||
}
|
||||
|
||||
output "pubsub_topic_short" {
|
||||
description = "Short Pub/Sub topic name."
|
||||
value = google_pubsub_topic.trigger_test.name
|
||||
}
|
||||
|
||||
output "eventarc_topic" {
|
||||
description = "Fully qualified Pub/Sub topic name that fires the Eventarc trigger."
|
||||
value = google_pubsub_topic.eventarc_source.id
|
||||
}
|
||||
|
||||
output "eventarc_topic_short" {
|
||||
description = "Short Eventarc source topic name."
|
||||
value = google_pubsub_topic.eventarc_source.name
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
output "project_id" {
|
||||
description = "GCP project ID used for test resources."
|
||||
value = var.project_id
|
||||
}
|
||||
|
||||
output "region" {
|
||||
description = "GCP region used for test resources."
|
||||
value = var.region
|
||||
}
|
||||
|
||||
output "suffix" {
|
||||
description = "The random suffix used for resource naming."
|
||||
value = local.suffix
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pub/Sub topic + push subscription pointing to /trigger/pubsub
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
resource "google_pubsub_topic" "trigger_test" {
|
||||
name = "${local.name_prefix}-${local.suffix}"
|
||||
project = var.project_id
|
||||
|
||||
depends_on = [google_project_service.apis]
|
||||
}
|
||||
|
||||
resource "google_pubsub_subscription" "trigger_push" {
|
||||
name = "${local.name_prefix}-push-${local.suffix}"
|
||||
project = var.project_id
|
||||
topic = google_pubsub_topic.trigger_test.id
|
||||
|
||||
push_config {
|
||||
push_endpoint = "${data.google_cloud_run_v2_service.trigger_agent.uri}/apps/trigger_echo_agent/trigger/pubsub"
|
||||
|
||||
oidc_token {
|
||||
service_account_email = google_service_account.pubsub_invoker.email
|
||||
audience = data.google_cloud_run_v2_service.trigger_agent.uri
|
||||
}
|
||||
}
|
||||
|
||||
# Short ack deadline for faster test feedback.
|
||||
ack_deadline_seconds = 30
|
||||
|
||||
# Retry policy for failed pushes.
|
||||
retry_policy {
|
||||
minimum_backoff = "10s"
|
||||
maximum_backoff = "60s"
|
||||
}
|
||||
|
||||
# Let the subscription persist until Terraform destroys it.
|
||||
# (default retention of 604800s is fine for tests)
|
||||
expiration_policy {
|
||||
ttl = ""
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
variable "project_id" {
|
||||
description = "GCP project ID for test resources."
|
||||
type = string
|
||||
}
|
||||
|
||||
variable "region" {
|
||||
description = "GCP region for Cloud Run and related resources."
|
||||
type = string
|
||||
default = "us-central1"
|
||||
}
|
||||
|
||||
variable "simulate_429_count" {
|
||||
description = "Number of 429 errors the echo agent simulates before succeeding (0 = disabled)."
|
||||
type = number
|
||||
default = 0
|
||||
}
|
||||
|
||||
variable "service_name" {
|
||||
description = "Optional name of a pre-deployed Cloud Run service. If provided, Terraform will use it instead of generating a name from the suffix."
|
||||
type = string
|
||||
default = null
|
||||
}
|
||||
|
||||
variable "suffix" {
|
||||
description = "Optional suffix for resource naming. If not provided, a random one will be generated."
|
||||
type = string
|
||||
default = null
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install the local ADK wheel (built from the repo under test).
|
||||
# Install the .whl file directly so pip uses it instead of PyPI.
|
||||
# Dependencies are resolved from PyPI automatically.
|
||||
COPY wheels/ /tmp/wheels/
|
||||
SHELL ["/bin/bash", "-c"]
|
||||
RUN pip install --no-cache-dir /tmp/wheels/google_adk-*.whl \
|
||||
&& rm -rf /tmp/wheels
|
||||
|
||||
# Copy the agent into the expected adk layout:
|
||||
# /app/agents/trigger_echo_agent/__init__.py
|
||||
# /app/agents/trigger_echo_agent/agent.py
|
||||
COPY __init__.py agents/trigger_echo_agent/__init__.py
|
||||
COPY agent.py agents/trigger_echo_agent/agent.py
|
||||
|
||||
ENV PORT=8080
|
||||
ENV HOST=0.0.0.0
|
||||
|
||||
EXPOSE 8080
|
||||
|
||||
CMD ["adk", "api_server", "--host=0.0.0.0", "--port=8080", "--trigger_sources=pubsub,eventarc", "/app/agents"]
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright 2026 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.
|
||||
@@ -0,0 +1,86 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Echo agent for remote trigger integration tests.
|
||||
|
||||
Uses a before_model_callback to echo user input without calling the LLM,
|
||||
making tests fast, deterministic, and free of LLM model quota usage.
|
||||
|
||||
Supports optional 429 simulation via the SIMULATE_429_COUNT environment
|
||||
variable: when set to N > 0, the first N invocations per session will raise
|
||||
a RuntimeError containing "429 RESOURCE_EXHAUSTED" before succeeding.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
from google.adk.agents.callback_context import CallbackContext
|
||||
from google.adk.agents.llm_agent import Agent
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.models.llm_response import LlmResponse
|
||||
from google.genai import types
|
||||
|
||||
# Track 429 simulation state across invocations within a process.
|
||||
# Keyed by session_id to allow per-session failure counts.
|
||||
_429_counter: dict[str, int] = {}
|
||||
|
||||
|
||||
def before_model_callback(
|
||||
callback_context: CallbackContext,
|
||||
llm_request: LlmRequest,
|
||||
) -> LlmResponse:
|
||||
"""Echo user input back without calling the LLM.
|
||||
|
||||
If SIMULATE_429_COUNT is set, raises a transient error for the first N
|
||||
invocations (per session) to exercise the retry-with-backoff logic in
|
||||
the trigger endpoints.
|
||||
"""
|
||||
fail_count = int(os.environ.get("SIMULATE_429_COUNT", "0"))
|
||||
session = callback_context.session
|
||||
session_id = session.id if session else "default"
|
||||
|
||||
if fail_count > 0:
|
||||
current = _429_counter.get(session_id, 0)
|
||||
if current < fail_count:
|
||||
_429_counter[session_id] = current + 1
|
||||
raise RuntimeError("429 RESOURCE_EXHAUSTED: simulated quota exceeded")
|
||||
|
||||
# Extract the most recent user message from the session events.
|
||||
user_text = ""
|
||||
if session and session.events:
|
||||
for event in reversed(session.events):
|
||||
if event.content and event.content.role == "user" and event.content.parts:
|
||||
user_text = event.content.parts[0].text or ""
|
||||
break
|
||||
|
||||
# Fall back to the current LLM request contents if no session event found.
|
||||
if not user_text and llm_request.contents:
|
||||
for content in reversed(llm_request.contents):
|
||||
if content.role == "user" and content.parts:
|
||||
user_text = content.parts[0].text or ""
|
||||
break
|
||||
|
||||
return LlmResponse(
|
||||
content=types.Content(
|
||||
role="model",
|
||||
parts=[types.Part(text=f"ECHO: {user_text}")],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
root_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="trigger_echo_agent",
|
||||
instruction="Echo agent for trigger testing.",
|
||||
before_model_callback=before_model_callback,
|
||||
)
|
||||
@@ -0,0 +1,125 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Remote integration tests for the /apps/trigger_echo_agent/trigger/eventarc endpoint.
|
||||
|
||||
Tests cover:
|
||||
|
||||
/apps/trigger_echo_agent/trigger/eventarc on Cloud Run.
|
||||
2. Full pipeline tests — publish to the Eventarc source topic and
|
||||
verify the request reached Cloud Run via Cloud Logging.
|
||||
|
||||
Prerequisites:
|
||||
- GCP project with Eventarc, Pub/Sub, Cloud Run, and Cloud Logging
|
||||
APIs enabled.
|
||||
- ``gcloud`` CLI authenticated.
|
||||
- Terraform >= 1.5 installed.
|
||||
|
||||
Run:
|
||||
GCP_PROJECT_ID=my-project pytest tests/remote/test_trigger_eventarc.py -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from google.cloud import logging as cloud_logging
|
||||
from google.cloud import pubsub_v1
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full Eventarc pipeline tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEventarcPipeline:
|
||||
"""Test the full Pub/Sub → Eventarc → Cloud Run pipeline.
|
||||
|
||||
Verification is done by checking Cloud Run request logs for successful
|
||||
POST requests to /apps/trigger_echo_agent/trigger/eventarc.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _setup(self, eventarc_topic, project_id):
|
||||
self.publisher = pubsub_v1.PublisherClient()
|
||||
self.topic_path = eventarc_topic
|
||||
self.project_id = project_id
|
||||
self.logging_client = cloud_logging.Client(project=project_id)
|
||||
|
||||
def _publish_event(self, data: str, wait_seconds: int = 45) -> None:
|
||||
"""Publish a message to the Eventarc source topic and wait."""
|
||||
future = self.publisher.publish(
|
||||
self.topic_path,
|
||||
data.encode("utf-8"),
|
||||
)
|
||||
future.result(timeout=30)
|
||||
time.sleep(wait_seconds)
|
||||
|
||||
def _count_successful_requests(
|
||||
self,
|
||||
path: str,
|
||||
since: datetime.datetime,
|
||||
) -> int:
|
||||
"""Count successful HTTP requests to the given path in Cloud Run logs."""
|
||||
timestamp = since.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
filter_str = (
|
||||
'resource.type="cloud_run_revision" '
|
||||
f'httpRequest.requestUrl:"{path}" '
|
||||
"httpRequest.status=200 "
|
||||
f'timestamp>="{timestamp}"'
|
||||
)
|
||||
entries = list(
|
||||
self.logging_client.list_entries(
|
||||
filter_=filter_str,
|
||||
page_size=50,
|
||||
)
|
||||
)
|
||||
return len(entries)
|
||||
|
||||
def test_eventarc_pubsub_trigger(self):
|
||||
"""Publish to Eventarc source topic and verify Cloud Run processes it."""
|
||||
start_time = datetime.datetime.now(datetime.timezone.utc)
|
||||
self._publish_event(json.dumps({"test": "eventarc-pipeline"}))
|
||||
count = self._count_successful_requests(
|
||||
"/apps/trigger_echo_agent/trigger/eventarc", start_time
|
||||
)
|
||||
assert count >= 1, (
|
||||
"Expected at least 1 successful request to"
|
||||
f" /apps/trigger_echo_agent/trigger/eventarc, found {count}"
|
||||
)
|
||||
|
||||
def test_eventarc_multiple_events(self):
|
||||
"""Publish multiple events and verify they are processed."""
|
||||
start_time = datetime.datetime.now(datetime.timezone.utc)
|
||||
for i in range(3):
|
||||
future = self.publisher.publish(
|
||||
self.topic_path,
|
||||
json.dumps({"seq": i}).encode("utf-8"),
|
||||
)
|
||||
future.result(timeout=30)
|
||||
|
||||
# Eventarc delivery + log ingestion can be slow; wait longer.
|
||||
time.sleep(90)
|
||||
|
||||
count = self._count_successful_requests(
|
||||
"/apps/trigger_echo_agent/trigger/eventarc", start_time
|
||||
)
|
||||
assert count >= 1, (
|
||||
"Expected at least 1 successful request to"
|
||||
f" /apps/trigger_echo_agent/trigger/eventarc, found {count}"
|
||||
)
|
||||
@@ -0,0 +1,166 @@
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Remote integration tests for the /apps/trigger_echo_agent/trigger/pubsub endpoint.
|
||||
|
||||
Tests cover:
|
||||
|
||||
2. Full pipeline tests — publish to a Pub/Sub topic with a push
|
||||
subscription pointing at the Cloud Run service, then verify via
|
||||
Cloud Logging that the requests reached the service.
|
||||
|
||||
Prerequisites:
|
||||
- GCP project with Pub/Sub, Cloud Run, and Cloud Logging APIs enabled.
|
||||
- ``gcloud`` CLI authenticated.
|
||||
- Terraform >= 1.5 installed.
|
||||
|
||||
Run:
|
||||
GCP_PROJECT_ID=my-project pytest tests/remote/test_trigger_pubsub.py -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import datetime
|
||||
import json
|
||||
import time
|
||||
|
||||
from google.cloud import logging as cloud_logging
|
||||
from google.cloud import pubsub_v1
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full Pub/Sub pipeline tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPubSubPipeline:
|
||||
"""Test the full Pub/Sub → push subscription → Cloud Run pipeline.
|
||||
|
||||
Verification is done by checking Cloud Run request logs for successful
|
||||
POST requests to /apps/trigger_echo_agent/trigger/pubsub. We look for
|
||||
httpRequest entries with
|
||||
status 200 that arrived after we published.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _setup(self, pubsub_topic, project_id, cloud_run_url):
|
||||
self.publisher = pubsub_v1.PublisherClient()
|
||||
self.topic_path = pubsub_topic
|
||||
self.project_id = project_id
|
||||
self.cloud_run_url = cloud_run_url
|
||||
self.logging_client = cloud_logging.Client(project=project_id)
|
||||
|
||||
def _publish_and_wait(
|
||||
self,
|
||||
data: str,
|
||||
attributes: dict[str, str] | None = None,
|
||||
wait_seconds: int = 30,
|
||||
) -> None:
|
||||
"""Publish a message and wait for it to be delivered."""
|
||||
future = self.publisher.publish(
|
||||
self.topic_path,
|
||||
data.encode("utf-8"),
|
||||
**(attributes or {}),
|
||||
)
|
||||
future.result(timeout=30)
|
||||
time.sleep(wait_seconds)
|
||||
|
||||
def _count_successful_requests(
|
||||
self,
|
||||
path: str,
|
||||
since: datetime.datetime,
|
||||
) -> int:
|
||||
"""Count successful HTTP requests to the given path in Cloud Run logs."""
|
||||
timestamp = since.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
filter_str = (
|
||||
'resource.type="cloud_run_revision" '
|
||||
f'httpRequest.requestUrl:"{path}" '
|
||||
"httpRequest.status=200 "
|
||||
f'timestamp>="{timestamp}"'
|
||||
)
|
||||
entries = list(
|
||||
self.logging_client.list_entries(
|
||||
filter_=filter_str,
|
||||
page_size=200,
|
||||
)
|
||||
)
|
||||
return len(entries)
|
||||
|
||||
def test_publish_text_message(self):
|
||||
"""Publish a text message and verify it reaches Cloud Run."""
|
||||
start_time = datetime.datetime.now(datetime.timezone.utc)
|
||||
self._publish_and_wait("hello-pipeline-test")
|
||||
count = self._count_successful_requests(
|
||||
"/apps/trigger_echo_agent/trigger/pubsub", start_time
|
||||
)
|
||||
assert count >= 1, (
|
||||
"Expected at least 1 successful request to"
|
||||
f" /apps/trigger_echo_agent/trigger/pubsub, found {count}"
|
||||
)
|
||||
|
||||
def test_publish_json_message(self):
|
||||
"""Publish a JSON message and verify processing."""
|
||||
start_time = datetime.datetime.now(datetime.timezone.utc)
|
||||
data = json.dumps({"type": "json", "value": 42})
|
||||
self._publish_and_wait(data)
|
||||
count = self._count_successful_requests(
|
||||
"/apps/trigger_echo_agent/trigger/pubsub", start_time
|
||||
)
|
||||
assert count >= 1, (
|
||||
"Expected at least 1 successful request to"
|
||||
f" /apps/trigger_echo_agent/trigger/pubsub, found {count}"
|
||||
)
|
||||
|
||||
def test_publish_with_attributes(self):
|
||||
"""Publish a message with attributes and verify processing."""
|
||||
start_time = datetime.datetime.now(datetime.timezone.utc)
|
||||
self._publish_and_wait(
|
||||
"attr-pipeline-test",
|
||||
attributes={"test_attr": "pytest"},
|
||||
)
|
||||
count = self._count_successful_requests(
|
||||
"/apps/trigger_echo_agent/trigger/pubsub", start_time
|
||||
)
|
||||
assert count >= 1, (
|
||||
"Expected at least 1 successful request to"
|
||||
f" /apps/trigger_echo_agent/trigger/pubsub, found {count}"
|
||||
)
|
||||
|
||||
def test_high_volume(self):
|
||||
"""Publish 20 messages and verify they are processed."""
|
||||
start_time = datetime.datetime.now(datetime.timezone.utc)
|
||||
n = 20
|
||||
futures = []
|
||||
for i in range(n):
|
||||
future = self.publisher.publish(
|
||||
self.topic_path,
|
||||
f"volume-test-{i}".encode("utf-8"),
|
||||
)
|
||||
futures.append(future)
|
||||
|
||||
for f in futures:
|
||||
f.result(timeout=30)
|
||||
|
||||
# Wait for push delivery.
|
||||
time.sleep(15)
|
||||
|
||||
count = self._count_successful_requests(
|
||||
"/apps/trigger_echo_agent/trigger/pubsub", start_time
|
||||
)
|
||||
# Allow some tolerance — at least 80% should arrive.
|
||||
assert (
|
||||
count >= n * 0.8
|
||||
), f"Expected at least {int(n * 0.8)} successful requests, found {count}"
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright 2026 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.
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright 2026 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.
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user