Skip to content

Commit 0b56f22

Browse files
google-genai-botcopybara-github
authored andcommitted
refactor(bigquery): Remove dataset_id from logging schema
Removes the dataset_id field from the BigQuery table schema and from each log entry created by the BigQueryAgentAnalyticsPlugin. This field is redundant, as all rows logged to a specific table belong to the same dataset. To ensure the plugin can still target the correct dataset, dataset_id is now a required argument in the BigQueryAgentAnalyticsPlugin constructor, and its default value has been removed. The BigQuery client user_agent is also updated with plugin version info to help identify traffic originating from this plugin. Unit tests have been updated to reflect the removal of dataset_id from log entries. PiperOrigin-RevId: 826596499
1 parent 156d235 commit 0b56f22

2 files changed

Lines changed: 2 additions & 17 deletions

File tree

src/google/adk/plugins/bigquery_logging_plugin.py

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
114114
def __init__(
115115
self,
116116
project_id: str,
117-
dataset_id: str = "adk_agent_logs",
117+
dataset_id: str,
118118
table_id: str = "agent_events",
119119
**kwargs,
120120
):
@@ -141,7 +141,7 @@ def _ensure_initialized_sync(self):
141141
scopes=["https://www.googleapis.com/auth/bigquery"]
142142
)
143143
client_info = google.api_core.client_info.ClientInfo(
144-
user_agent=f"google-adk-plugin/{version.__version__}"
144+
user_agent=f"google-adk-bq-logger/{version.__version__}"
145145
)
146146
self._bq_client = bigquery.Client(
147147
project=self._project_id,
@@ -157,7 +157,6 @@ def _ensure_initialized_sync(self):
157157
table_ref = dataset_ref.table(self._table_id)
158158
# Schema without separate token columns
159159
schema = [
160-
bigquery.SchemaField("dataset_id", "STRING"),
161160
bigquery.SchemaField("timestamp", "TIMESTAMP"),
162161
bigquery.SchemaField("event_type", "STRING"),
163162
bigquery.SchemaField("agent", "STRING"),
@@ -189,7 +188,6 @@ def _sync_log():
189188
self._table_id
190189
)
191190
default_row = {
192-
"dataset_id": None,
193191
"timestamp": datetime.now(timezone.utc).isoformat(),
194192
"event_type": None,
195193
"agent": None,
@@ -226,7 +224,6 @@ async def on_user_message_callback(
226224
) -> Optional[types.Content]:
227225
"""Log user message and invocation start."""
228226
event_dict = {
229-
"dataset_id": self._dataset_id,
230227
"timestamp": datetime.now(timezone.utc).isoformat(),
231228
"event_type": "USER_MESSAGE_RECEIVED",
232229
"agent": invocation_context.agent.name,
@@ -243,7 +240,6 @@ async def before_run_callback(
243240
) -> Optional[types.Content]:
244241
"""Log invocation start."""
245242
event_dict = {
246-
"dataset_id": self._dataset_id,
247243
"timestamp": datetime.now(timezone.utc).isoformat(),
248244
"event_type": "INVOCATION_STARTING",
249245
"agent": invocation_context.agent.name,
@@ -259,7 +255,6 @@ async def on_event_callback(
259255
) -> Optional[Event]:
260256
"""Logs event data to BigQuery."""
261257
event_dict = {
262-
"dataset_id": self._dataset_id,
263258
"timestamp": datetime.fromtimestamp(
264259
event.timestamp, timezone.utc
265260
).isoformat(),
@@ -285,7 +280,6 @@ async def after_run_callback(
285280
) -> Optional[None]:
286281
"""Log invocation completion."""
287282
event_dict = {
288-
"dataset_id": self._dataset_id,
289283
"timestamp": datetime.now(timezone.utc).isoformat(),
290284
"event_type": "INVOCATION_COMPLETED",
291285
"agent": invocation_context.agent.name,
@@ -301,7 +295,6 @@ async def before_agent_callback(
301295
) -> Optional[types.Content]:
302296
"""Log agent execution start."""
303297
event_dict = {
304-
"dataset_id": self._dataset_id,
305298
"timestamp": datetime.now(timezone.utc).isoformat(),
306299
"event_type": "AGENT_STARTING",
307300
"agent": agent.name,
@@ -318,7 +311,6 @@ async def after_agent_callback(
318311
) -> Optional[types.Content]:
319312
"""Log agent execution completion."""
320313
event_dict = {
321-
"dataset_id": self._dataset_id,
322314
"timestamp": datetime.now(timezone.utc).isoformat(),
323315
"event_type": "AGENT_COMPLETED",
324316
"agent": agent.name,
@@ -386,7 +378,6 @@ async def before_model_callback(
386378
final_content = " | ".join(content_parts)
387379

388380
event_dict = {
389-
"dataset_id": self._dataset_id,
390381
"timestamp": datetime.now(timezone.utc).isoformat(),
391382
"event_type": "LLM_REQUEST",
392383
"agent": callback_context.agent_name,
@@ -444,7 +435,6 @@ async def after_model_callback(
444435
final_content = " | ".join(content_parts)
445436

446437
event_dict = {
447-
"dataset_id": self._dataset_id,
448438
"timestamp": datetime.now(timezone.utc).isoformat(),
449439
"event_type": "LLM_RESPONSE",
450440
"agent": callback_context.agent_name,
@@ -468,7 +458,6 @@ async def before_tool_callback(
468458
) -> Optional[None]:
469459
"""Log tool execution start."""
470460
event_dict = {
471-
"dataset_id": self._dataset_id,
472461
"timestamp": datetime.now(timezone.utc).isoformat(),
473462
"event_type": "TOOL_STARTING",
474463
"agent": tool_context.agent_name,
@@ -493,7 +482,6 @@ async def after_tool_callback(
493482
) -> None:
494483
"""Log tool execution completion."""
495484
event_dict = {
496-
"dataset_id": self._dataset_id,
497485
"timestamp": datetime.now(timezone.utc).isoformat(),
498486
"event_type": "TOOL_COMPLETED",
499487
"agent": tool_context.agent_name,
@@ -514,7 +502,6 @@ async def on_model_error_callback(
514502
) -> Optional[LlmResponse]:
515503
"""Log LLM error."""
516504
event_dict = {
517-
"dataset_id": self._dataset_id,
518505
"timestamp": datetime.now(timezone.utc).isoformat(),
519506
"event_type": "LLM_ERROR",
520507
"agent": callback_context.agent_name,
@@ -536,7 +523,6 @@ async def on_tool_error_callback(
536523
) -> None:
537524
"""Log tool error."""
538525
event_dict = {
539-
"dataset_id": self._dataset_id,
540526
"timestamp": datetime.now(timezone.utc).isoformat(),
541527
"event_type": "TOOL_ERROR",
542528
"agent": tool_context.agent_name,

tests/unittests/plugins/test_bigquery_logging_plugin.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,6 @@ def _get_logged_entry(self):
127127
return rows[0]
128128

129129
def _assert_common_fields(self, log_entry, event_type):
130-
assert log_entry["dataset_id"] == self.dataset_id
131130
assert log_entry["event_type"] == event_type
132131
assert log_entry["agent"] == "MyTestAgent"
133132
assert log_entry["session_id"] == "session-123"

0 commit comments

Comments
 (0)