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

PiperOrigin-RevId: 896612027

Author: google-genai-bot

src/google/adk/agents/config_agent_utils.py

@@ -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,6 +55,7 @@
 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 ..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 @@ async def get_runner_async(self, app_name: str) -> Runner:
     # 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 @@ async def patch_memory(
         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 @@ async def run_agent(req: RunAgentRequest) -> list[Event]:
     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 @@ async def run_agent_live(
         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

@@ -67,6 +67,7 @@
 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 @@ def __init__(
     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 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None:
       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 @@ def __init__(
     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 @@ def get_credentials():
         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

@@ -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 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

@@ -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
+)