github-featured/google--adk-python
2
0
Fork 0
Code Issues Pull Requests Actions 7 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

feat: Support loading agents from Visual Builder with BigQuery-powered logging

Browse Source 
PiperOrigin-RevId: 896612027
This commit is contained in:
Google Team Member
2026-04-08 11:33:28 -07:00
committed by Copybara-Service
parent 9e73ab8466
commit 20748894cd
8 changed files with 509 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/agents/config_agent_utils.py
+1 -1
View File
@@ -36,7 +36,7 @@ def from_config(config_path: str) -> BaseAgent:
"""Build agent from a configfile path.
Args:
config: the path to a YAML config file.
config_path: the path to a YAML config file.
Returns:
The created agent instance.
src/google/adk/cli/adk_web_server.py
+55 -1
View File
@@ -55,6 +55,7 @@ from starlette.types import Lifespan
from typing_extensions import deprecated
from typing_extensions import override
from watchdog.observers import Observer
import yaml
from . import agent_graph
from ..agents.base_agent import BaseAgent
@@ -89,6 +90,7 @@ from ..evaluation.eval_sets_manager import EvalSetsManager
from ..events.event import Event
from ..memory.base_memory_service import BaseMemoryService
from ..plugins.base_plugin import BasePlugin
from ..plugins.bigquery_agent_analytics_plugin import BigQueryAgentAnalyticsPlugin
from ..runners import Runner
from ..sessions.base_session_service import BaseSessionService
from ..sessions.session import Session
@@ -697,17 +699,55 @@ class AdkWebServer:
# Instantiate extra plugins if configured
extra_plugins_instances = self._instantiate_extra_plugins()
plugins_yaml_path = os.path.join(self.agents_dir, app_name, "plugins.yaml")
bq_analytics_config = None
if os.path.exists(plugins_yaml_path):
with open(plugins_yaml_path, "r", encoding="utf-8") as f:
plugins_config = yaml.safe_load(f)
if plugins_config and isinstance(plugins_config, dict):
bq_analytics_config = plugins_config.get("bigquery_agent_analytics")
# Determine if the agent was loaded from YAML based on the agent loader info
is_visual_builder = False
detailed_agents = self.agent_loader.list_agents_detailed()
for agent_info in detailed_agents:
if agent_info.get("name") == app_name:
if agent_info.get("language") == "yaml":
is_visual_builder = True
break
if isinstance(agent_or_app, BaseAgent):
plugins = extra_plugins_instances
# Handle BigQuery Analytics Plugin injection
if bq_analytics_config and all([
bq_analytics_config.get("project_id"),
bq_analytics_config.get("dataset_id"),
bq_analytics_config.get("dataset_location"),
]):
plugins.append(
BigQueryAgentAnalyticsPlugin(
project_id=bq_analytics_config.get("project_id"),
dataset_id=bq_analytics_config.get("dataset_id"),
table_id=bq_analytics_config.get("table_id"),
location=bq_analytics_config.get("dataset_location"),
)
)
agentic_app = App(
name=app_name,
root_agent=agent_or_app,
plugins=extra_plugins_instances,
plugins=plugins,
)
else:
# Combine existing plugins with extra plugins
agent_or_app.plugins = agent_or_app.plugins + extra_plugins_instances
agentic_app = agent_or_app
# If the root agent was loaded from YAML, we treat it as being from Visual Builder
if is_visual_builder:
object.__setattr__(agentic_app, "_is_visual_builder_app", True)
runner = self._create_runner(agentic_app)
self.runner_dict[app_name] = runner
return runner
@@ -1840,9 +1880,20 @@ class AdkWebServer:
raise HTTPException(status_code=404, detail="Session not found")
await self.memory_service.add_session_to_memory(session)
def _set_telemetry_context_if_needed(runner: Runner):
"""Helper to set contextvars for the current request task."""
app = getattr(runner, "app", None)
from ..utils._telemetry_context import _is_visual_builder
if app and getattr(app, "_is_visual_builder_app", False):
_is_visual_builder.set(True)
else:
_is_visual_builder.set(False)
@app.post("/run", response_model_exclude_none=True)
async def run_agent(req: RunAgentRequest) -> list[Event]:
runner = await self.get_runner_async(req.app_name)
_set_telemetry_context_if_needed(runner)
try:
async with Aclosing(
runner.run_async(
@@ -1864,6 +1915,7 @@ class AdkWebServer:
async def run_agent_sse(req: RunAgentRequest) -> StreamingResponse:
stream_mode = StreamingMode.SSE if req.streaming else StreamingMode.NONE
runner = await self.get_runner_async(req.app_name)
_set_telemetry_context_if_needed(runner)
# Validate session existence before starting the stream.
# We check directly here instead of eagerly advancing the
@@ -2039,6 +2091,8 @@ class AdkWebServer:
return
await websocket.accept()
runner_for_context = await self.get_runner_async(app_name)
_set_telemetry_context_if_needed(runner_for_context)
session = await self.session_service.get_session(
app_name=app_name, user_id=user_id, session_id=session_id
src/google/adk/plugins/bigquery_agent_analytics_plugin.py
+19 -4
View File
@@ -67,6 +67,7 @@ from ..models.llm_request import LlmRequest
from ..models.llm_response import LlmResponse
from ..tools.base_tool import BaseTool
from ..tools.tool_context import ToolContext
from ..utils._telemetry_context import _is_visual_builder
from ..version import __version__
from .base_plugin import BasePlugin
@@ -922,6 +923,9 @@ class BatchProcessor:
self.flush_interval = flush_interval
self.retry_config = retry_config
self.shutdown_timeout = shutdown_timeout
self._visual_builder = _is_visual_builder.get()
self._queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(
maxsize=queue_max_size
)
@@ -1092,9 +1096,15 @@ class BatchProcessor:
serialized_schema = self.arrow_schema.serialize().to_pybytes()
serialized_batch = arrow_batch.serialize().to_pybytes()
trace_id_prefix = (
"google-adk-bq-logger-visual-builder"
if self._visual_builder
else "google-adk-bq-logger"
)
req = bq_storage_types.AppendRowsRequest(
write_stream=self.write_stream,
trace_id=f"google-adk-bq-logger/{__version__}",
trace_id=f"{trace_id_prefix}/{__version__}",
)
req.arrow_rows.writer_schema.serialized_schema = serialized_schema
req.arrow_rows.rows.serialized_record_batch = serialized_batch
@@ -1900,6 +1910,8 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
self.table_id = table_id or self.config.table_id
self.location = location
self._visual_builder = _is_visual_builder.get()
self._started = False
self._startup_error: Optional[Exception] = None
self._is_shutting_down = False
@@ -2030,9 +2042,12 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
if quota_project_id
else None
)
client_info = gapic_client_info.ClientInfo(
user_agent=f"google-adk-bq-logger/{__version__}"
)
user_agents = [f"google-adk-bq-logger/{__version__}"]
if self._visual_builder:
user_agents.append(f"google-adk-visual-builder/{__version__}")
client_info = gapic_client_info.ClientInfo(user_agent=" ".join(user_agents))
write_client = BigQueryWriteAsyncClient(
credentials=creds,
src/google/adk/tools/bigquery/client.py
+13
View File
@@ -14,6 +14,7 @@
from __future__ import annotations
import os
from typing import List
from typing import Optional
from typing import Union
@@ -25,12 +26,16 @@ from google.cloud import bigquery
from google.cloud import dataplex_v1
from ... import version
from ...utils._telemetry_context import _is_visual_builder
USER_AGENT_BASE = f"google-adk/{version.__version__}"
BQ_USER_AGENT = f"adk-bigquery-tool {USER_AGENT_BASE}"
DP_USER_AGENT = f"adk-dataplex-tool {USER_AGENT_BASE}"
USER_AGENT = BQ_USER_AGENT
# Internal identifier for Visual Builder usage tracking.
_VISUAL_BUILDER_UA = "google-adk-visual-builder"
def get_bigquery_client(
*,
@@ -52,6 +57,10 @@ def get_bigquery_client(
"""
user_agents = [BQ_USER_AGENT]
if _is_visual_builder.get():
user_agents.append(_VISUAL_BUILDER_UA)
if user_agent:
if isinstance(user_agent, str):
user_agents.append(user_agent)
@@ -88,6 +97,10 @@ def get_dataplex_catalog_client(
"""
user_agents = [DP_USER_AGENT]
if _is_visual_builder.get():
user_agents.append(_VISUAL_BUILDER_UA)
if user_agent:
if isinstance(user_agent, str):
user_agents.append(user_agent)
src/google/adk/utils/_telemetry_context.py
+25
View File
@@ -0,0 +1,25 @@
# 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.
"""Context variables for internal telemetry use."""
from __future__ import annotations
import contextvars
# Internal context variable for Visual Builder usage tracking.
# True if the current execution is within a Visual Builder context.
_is_visual_builder: contextvars.ContextVar[bool] = contextvars.ContextVar(
"_is_visual_builder", default=False
)
tests/unittests/cli/test_fast_api.py
+253 -8
View File
@@ -40,6 +40,7 @@ from google.adk.evaluation.eval_result import EvalSetResult
from google.adk.evaluation.in_memory_eval_sets_manager import InMemoryEvalSetsManager
from google.adk.events.event import Event
from google.adk.events.event_actions import EventActions
from google.adk.plugins.bigquery_agent_analytics_plugin import BigQueryAgentAnalyticsPlugin
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.sessions.session import Session
@@ -187,19 +188,39 @@ def mock_agent_loader():
pass
def load_agent(self, app_name):
if app_name == "yaml_app" or app_name == "bq_app":
agent = DummyAgent(name="yaml_agent")
agent._config = MagicMock(logging=None)
return agent
return root_agent
def list_agents(self):
return ["test_app"]
return ["test_app", "yaml_app", "bq_app"]
def list_agents_detailed(self):
return [{
"name": "test_app",
"root_agent_name": "test_agent",
"description": "A test agent for unit testing",
"language": "python",
"is_computer_use": False,
}]
return [
{
"name": "test_app",
"root_agent_name": "test_agent",
"description": "A test agent for unit testing",
"language": "python",
"is_computer_use": False,
},
{
"name": "yaml_app",
"root_agent_name": "yaml_agent",
"description": "A yaml agent for unit testing",
"language": "yaml",
"is_computer_use": False,
},
{
"name": "bq_app",
"root_agent_name": "yaml_agent",
"description": "A bq agent for unit testing",
"language": "yaml",
"is_computer_use": False,
},
]
return MockAgentLoader(".")
@@ -518,6 +539,103 @@ def _create_test_client(
return TestClient(app)
def test_agent_with_bigquery_analytics_plugin(
tmp_path,
mock_session_service,
mock_artifact_service,
mock_memory_service,
mock_agent_loader,
mock_eval_sets_manager,
mock_eval_set_results_manager,
):
"""Verify that plugins.yaml is correctly read to attach BigQueryAgentAnalyticsPlugin."""
app_name = "bq_app"
app_dir = tmp_path / app_name
app_dir.mkdir(parents=True)
plugins_yaml_content = """\
bigquery_agent_analytics:
project_id: test-project
dataset_id: test-dataset
table_id: test-table
dataset_location: US
"""
(app_dir / "plugins.yaml").write_text(plugins_yaml_content)
with (
patch.object(signal, "signal", autospec=True, return_value=None),
patch.object(
fast_api_module,
"create_session_service_from_options",
autospec=True,
return_value=mock_session_service,
),
patch.object(
fast_api_module,
"create_artifact_service_from_options",
autospec=True,
return_value=mock_artifact_service,
),
patch.object(
fast_api_module,
"create_memory_service_from_options",
autospec=True,
return_value=mock_memory_service,
),
patch.object(
fast_api_module,
"AgentLoader",
autospec=True,
return_value=mock_agent_loader,
),
patch.object(
fast_api_module,
"LocalEvalSetsManager",
autospec=True,
return_value=mock_eval_sets_manager,
),
patch.object(
fast_api_module,
"LocalEvalSetResultsManager",
autospec=True,
return_value=mock_eval_set_results_manager,
),
):
from google.adk.cli.adk_web_server import AdkWebServer
adk_web_server = AdkWebServer(
agent_loader=mock_agent_loader,
session_service=mock_session_service,
memory_service=mock_memory_service,
artifact_service=mock_artifact_service,
credential_service=MagicMock(),
eval_sets_manager=mock_eval_sets_manager,
eval_set_results_manager=mock_eval_set_results_manager,
agents_dir=str(tmp_path),
)
runner = asyncio.run(adk_web_server.get_runner_async(app_name))
# Assert that the plugin was attached
assert any(
isinstance(p, BigQueryAgentAnalyticsPlugin) for p in runner.app.plugins
)
# Check the configuration of the plugin
bq_plugin = next(
p
for p in runner.app.plugins
if isinstance(p, BigQueryAgentAnalyticsPlugin)
)
assert bq_plugin.project_id == "test-project"
assert bq_plugin.dataset_id == "test-dataset"
assert bq_plugin.table_id == "test-table"
assert bq_plugin.location == "US"
# Assert that the internal visual builder flag is set on the app
assert getattr(runner.app, "_is_visual_builder_app", False) is True
@pytest.fixture
def test_app(
mock_session_service,
@@ -2178,5 +2296,132 @@ def test_returns_404_without_auto_create(
assert "Session not found" in response.json()["detail"]
@pytest.mark.asyncio
async def test_independent_telemetry_context(
mock_session_service,
mock_artifact_service,
mock_memory_service,
mock_agent_loader,
mock_eval_sets_manager,
mock_eval_set_results_manager,
monkeypatch,
):
"""Test that two agents have independent is_visual_builder context variables."""
from google.adk.utils._telemetry_context import _is_visual_builder
import httpx
# We use httpx.AsyncClient to send concurrent requests to the FastAPI app.
# This proves that is_visual_builder doesn't leak across concurrent requests.
captured_visual_builder_values = {}
async def run_async_capture(
self,
*,
user_id: str,
session_id: str,
invocation_id: Optional[str] = None,
new_message: Optional[types.Content] = None,
state_delta: Optional[dict[str, Any]] = None,
run_config: Optional[RunConfig] = None,
):
# Capture the value of is_visual_builder inside the request context
captured_visual_builder_values[self.app.name] = _is_visual_builder.get()
# Sleep to ensure both requests overlap in time
await asyncio.sleep(0.1)
# Read again to ensure it wasn't overwritten by the other concurrent request
captured_visual_builder_values[self.app.name + "_after_sleep"] = (
_is_visual_builder.get()
)
yield _event_1()
monkeypatch.setattr(Runner, "run_async", run_async_capture)
with (
patch.object(signal, "signal", autospec=True, return_value=None),
patch.object(
fast_api_module,
"create_session_service_from_options",
autospec=True,
return_value=mock_session_service,
),
patch.object(
fast_api_module,
"create_artifact_service_from_options",
autospec=True,
return_value=mock_artifact_service,
),
patch.object(
fast_api_module,
"create_memory_service_from_options",
autospec=True,
return_value=mock_memory_service,
),
patch.object(
fast_api_module,
"AgentLoader",
autospec=True,
return_value=mock_agent_loader,
),
patch.object(
fast_api_module,
"LocalEvalSetsManager",
autospec=True,
return_value=mock_eval_sets_manager,
),
patch.object(
fast_api_module,
"LocalEvalSetResultsManager",
autospec=True,
return_value=mock_eval_set_results_manager,
),
):
app = get_fast_api_app(
agents_dir=".",
web=True,
session_service_uri="",
artifact_service_uri="",
memory_service_uri="",
allow_origins=["*"],
a2a=False,
host="127.0.0.1",
port=8000,
)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(
transport=transport, base_url="http://test"
) as client:
# Send concurrent requests
req1 = client.post(
"/run",
json={
"app_name": "test_app",
"user_id": "test_user",
"session_id": "test_session",
"new_message": {"role": "user", "parts": [{"text": "Hello"}]},
},
)
req2 = client.post(
"/run",
json={
"app_name": "yaml_app",
"user_id": "test_user",
"session_id": "test_session",
"new_message": {"role": "user", "parts": [{"text": "Hello"}]},
},
)
await asyncio.gather(req1, req2)
assert captured_visual_builder_values.get("test_app") == False
assert captured_visual_builder_values.get("test_app_after_sleep") == False
assert captured_visual_builder_values.get("yaml_app") == True
assert captured_visual_builder_values.get("yaml_app_after_sleep") == True
if __name__ == "__main__":
pytest.main(["-xvs", __file__])
tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py
+115 -1
View File
@@ -33,6 +33,7 @@ from google.adk.sessions import base_session_service as base_session_service_lib
from google.adk.sessions import session as session_lib
from google.adk.tools import base_tool as base_tool_lib
from google.adk.tools import tool_context as tool_context_lib
from google.adk.utils._telemetry_context import _is_visual_builder
from google.adk.version import __version__
import google.auth
from google.auth import exceptions as auth_exceptions
@@ -279,7 +280,8 @@ async def _get_captured_event_dict_async(mock_write_client, expected_schema):
assert len(requests) == 1
request = requests[0]
assert request.write_stream == DEFAULT_STREAM_NAME
assert request.trace_id == f"google-adk-bq-logger/{__version__}"
assert request.trace_id.startswith("google-adk-bq-logger")
assert request.trace_id.endswith(f"/{__version__}")
# Parse the Arrow batch back to a dict for verification
try:
reader = pa.ipc.open_stream(request.arrow_rows.rows.serialized_record_batch)
@@ -2246,6 +2248,118 @@ class TestBigQueryAgentAnalyticsPlugin:
assert format(finished_spans[0].context.span_id, "016x") == span_id
assert format(finished_spans[0].context.trace_id, "032x") == trace_id
@pytest.mark.asyncio
async def test_keyword_identifiers_emission_default(
self,
mock_auth_default,
mock_bq_client,
callback_context,
):
"""Verify the default keyword flow for User-Agent and Trace-ID."""
keyword = "google-adk-bq-logger"
mock_write_client = mock.AsyncMock()
# 1. Verify User-Agent contains default keyword.
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.BigQueryWriteAsyncClient",
autospec=True,
) as mock_write_cls:
mock_write_cls.return_value = mock_write_client
async with managed_plugin(PROJECT_ID, DATASET_ID) as plugin:
await plugin._ensure_started()
_, kwargs = mock_write_cls.call_args
client_info = kwargs.get("client_info")
assert f"{keyword}/{__version__}" in client_info.user_agent
# 2. Verify Trace ID contains default keyword.
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.BigQueryWriteAsyncClient",
autospec=True,
) as mock_write_cls:
mock_write_cls.return_value = mock_write_client
async with managed_plugin(PROJECT_ID, DATASET_ID) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Hi")])],
)
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await plugin.flush()
call_args = mock_write_client.append_rows.call_args
requests_iter = call_args.args[0]
requests = []
async for req in requests_iter:
requests.append(req)
assert requests[0].trace_id.startswith(keyword)
assert requests[0].trace_id.endswith(f"/{__version__}")
@pytest.mark.asyncio
async def test_visual_builder_identifiers_flow(
self,
mock_auth_default,
mock_bq_client,
callback_context,
dummy_arrow_schema,
):
"""Verify visual-builder keyword flow via contextvars."""
keyword = "google-adk-visual-builder"
mock_write_client = mock.AsyncMock()
# Simulate setting the internal flag via contextvars
token = _is_visual_builder.set(True)
try:
# 1. Verify Client User-Agent
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.BigQueryWriteAsyncClient",
autospec=True,
) as mock_write_cls:
mock_write_cls.return_value = mock_write_client
async with managed_plugin(PROJECT_ID, DATASET_ID) as plugin:
await plugin._ensure_started()
_, kwargs = mock_write_cls.call_args
client_info = kwargs.get("client_info")
assert keyword in client_info.user_agent
# 2. Verify Request Trace ID
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.BigQueryWriteAsyncClient",
autospec=True,
) as mock_write_cls:
mock_write_cls.return_value = mock_write_client
async with managed_plugin(PROJECT_ID, DATASET_ID) as plugin:
await plugin._ensure_started()
mock_write_client.append_rows.reset_mock()
llm_request = llm_request_lib.LlmRequest(
model="gemini-pro",
contents=[types.Content(parts=[types.Part(text="Hi")])],
)
await plugin.before_model_callback(
callback_context=callback_context, llm_request=llm_request
)
await plugin.flush()
call_args = mock_write_client.append_rows.call_args
requests_iter = call_args.args[0]
requests = []
async for req in requests_iter:
requests.append(req)
assert requests[0].trace_id.startswith(
"google-adk-bq-logger-visual-builder"
)
assert requests[0].trace_id.endswith(f"/{__version__}")
finally:
_is_visual_builder.reset(token)
@pytest.mark.asyncio
async def test_flush_mechanism(
self,
tests/unittests/tools/bigquery/test_bigquery_client.py
+28
View File
@@ -21,6 +21,7 @@ import google.adk
from google.adk.tools.bigquery.client import DP_USER_AGENT
from google.adk.tools.bigquery.client import get_bigquery_client
from google.adk.tools.bigquery.client import get_dataplex_catalog_client
from google.adk.utils._telemetry_context import _is_visual_builder
from google.api_core.gapic_v1 import client_info as gapic_client_info
import google.auth
from google.auth.exceptions import DefaultCredentialsError
@@ -193,6 +194,33 @@ def test_bigquery_client_user_agent_custom_list():
assert expected_user_agents.issubset(actual_user_agents)
def test_bigquery_client_user_agent_visual_builder():
"""Test BigQuery client user agent when visual builder flag is set."""
token = _is_visual_builder.set(True)
try:
with mock.patch.object(
bigquery_client, "Connection", autospec=True
) as mock_connection:
# Trigger the BigQuery client creation
get_bigquery_client(
project="test-gcp-project",
credentials=mock.create_autospec(Credentials, instance=True),
)
# Verify that the tracking user agent was set
client_info_arg = mock_connection.call_args[1].get("client_info")
assert client_info_arg is not None
expected_user_agents = {
"adk-bigquery-tool",
f"google-adk/{google.adk.__version__}",
"google-adk-visual-builder",
}
actual_user_agents = set(client_info_arg.user_agent.split())
assert expected_user_agents.issubset(actual_user_agents)
finally:
_is_visual_builder.reset(token)
def test_bigquery_client_location_custom():
"""Test BigQuery client custom location."""
# Trigger the BigQuery client creation