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
This commit is contained in:
Google Team Member
2025-10-31 12:57:54 -07:00
committed by Copybara-Service
parent 156d235479
commit 0b56f2287b
2 changed files with 2 additions and 17 deletions
@@ -114,7 +114,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
def __init__( def __init__(
self, self,
project_id: str, project_id: str,
dataset_id: str = "adk_agent_logs", dataset_id: str,
table_id: str = "agent_events", table_id: str = "agent_events",
**kwargs, **kwargs,
): ):
@@ -141,7 +141,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
scopes=["https://www.googleapis.com/auth/bigquery"] scopes=["https://www.googleapis.com/auth/bigquery"]
) )
client_info = google.api_core.client_info.ClientInfo( client_info = google.api_core.client_info.ClientInfo(
user_agent=f"google-adk-plugin/{version.__version__}" user_agent=f"google-adk-bq-logger/{version.__version__}"
) )
self._bq_client = bigquery.Client( self._bq_client = bigquery.Client(
project=self._project_id, project=self._project_id,
@@ -157,7 +157,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
table_ref = dataset_ref.table(self._table_id) table_ref = dataset_ref.table(self._table_id)
# Schema without separate token columns # Schema without separate token columns
schema = [ schema = [
bigquery.SchemaField("dataset_id", "STRING"),
bigquery.SchemaField("timestamp", "TIMESTAMP"), bigquery.SchemaField("timestamp", "TIMESTAMP"),
bigquery.SchemaField("event_type", "STRING"), bigquery.SchemaField("event_type", "STRING"),
bigquery.SchemaField("agent", "STRING"), bigquery.SchemaField("agent", "STRING"),
@@ -189,7 +188,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
self._table_id self._table_id
) )
default_row = { default_row = {
"dataset_id": None,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": None, "event_type": None,
"agent": None, "agent": None,
@@ -226,7 +224,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> Optional[types.Content]: ) -> Optional[types.Content]:
"""Log user message and invocation start.""" """Log user message and invocation start."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "USER_MESSAGE_RECEIVED", "event_type": "USER_MESSAGE_RECEIVED",
"agent": invocation_context.agent.name, "agent": invocation_context.agent.name,
@@ -243,7 +240,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> Optional[types.Content]: ) -> Optional[types.Content]:
"""Log invocation start.""" """Log invocation start."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "INVOCATION_STARTING", "event_type": "INVOCATION_STARTING",
"agent": invocation_context.agent.name, "agent": invocation_context.agent.name,
@@ -259,7 +255,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> Optional[Event]: ) -> Optional[Event]:
"""Logs event data to BigQuery.""" """Logs event data to BigQuery."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.fromtimestamp( "timestamp": datetime.fromtimestamp(
event.timestamp, timezone.utc event.timestamp, timezone.utc
).isoformat(), ).isoformat(),
@@ -285,7 +280,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> Optional[None]: ) -> Optional[None]:
"""Log invocation completion.""" """Log invocation completion."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "INVOCATION_COMPLETED", "event_type": "INVOCATION_COMPLETED",
"agent": invocation_context.agent.name, "agent": invocation_context.agent.name,
@@ -301,7 +295,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> Optional[types.Content]: ) -> Optional[types.Content]:
"""Log agent execution start.""" """Log agent execution start."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "AGENT_STARTING", "event_type": "AGENT_STARTING",
"agent": agent.name, "agent": agent.name,
@@ -318,7 +311,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> Optional[types.Content]: ) -> Optional[types.Content]:
"""Log agent execution completion.""" """Log agent execution completion."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "AGENT_COMPLETED", "event_type": "AGENT_COMPLETED",
"agent": agent.name, "agent": agent.name,
@@ -386,7 +378,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
final_content = " | ".join(content_parts) final_content = " | ".join(content_parts)
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "LLM_REQUEST", "event_type": "LLM_REQUEST",
"agent": callback_context.agent_name, "agent": callback_context.agent_name,
@@ -444,7 +435,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
final_content = " | ".join(content_parts) final_content = " | ".join(content_parts)
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "LLM_RESPONSE", "event_type": "LLM_RESPONSE",
"agent": callback_context.agent_name, "agent": callback_context.agent_name,
@@ -468,7 +458,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> Optional[None]: ) -> Optional[None]:
"""Log tool execution start.""" """Log tool execution start."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "TOOL_STARTING", "event_type": "TOOL_STARTING",
"agent": tool_context.agent_name, "agent": tool_context.agent_name,
@@ -493,7 +482,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> None: ) -> None:
"""Log tool execution completion.""" """Log tool execution completion."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "TOOL_COMPLETED", "event_type": "TOOL_COMPLETED",
"agent": tool_context.agent_name, "agent": tool_context.agent_name,
@@ -514,7 +502,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> Optional[LlmResponse]: ) -> Optional[LlmResponse]:
"""Log LLM error.""" """Log LLM error."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "LLM_ERROR", "event_type": "LLM_ERROR",
"agent": callback_context.agent_name, "agent": callback_context.agent_name,
@@ -536,7 +523,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
) -> None: ) -> None:
"""Log tool error.""" """Log tool error."""
event_dict = { event_dict = {
"dataset_id": self._dataset_id,
"timestamp": datetime.now(timezone.utc).isoformat(), "timestamp": datetime.now(timezone.utc).isoformat(),
"event_type": "TOOL_ERROR", "event_type": "TOOL_ERROR",
"agent": tool_context.agent_name, "agent": tool_context.agent_name,
@@ -127,7 +127,6 @@ class TestBigQueryAgentAnalyticsPlugin(PluginTestBase):
return rows[0] return rows[0]
def _assert_common_fields(self, log_entry, event_type): def _assert_common_fields(self, log_entry, event_type):
assert log_entry["dataset_id"] == self.dataset_id
assert log_entry["event_type"] == event_type assert log_entry["event_type"] == event_type
assert log_entry["agent"] == "MyTestAgent" assert log_entry["agent"] == "MyTestAgent"
assert log_entry["session_id"] == "session-123" assert log_entry["session_id"] == "session-123"