Skip to content

Commit 34fea3e

Browse files
committed
feat(bigquery): add default_attributes config to BigQueryAgentAnalyticsPlugin
Add a new `default_attributes` configuration option to BigQueryLoggerConfig that allows users to inject static key-value pairs into every logged event's attributes field. This is useful for adding deployment metadata like service version, environment, or commit SHA to all analytics events. Features: - New `default_attributes: Optional[dict[str, Any]]` config field - Default attributes are merged into every event's attributes - Event-specific attributes override defaults when there are conflicts - Fully backward compatible (None by default) Example usage: ```python config = BigQueryLoggerConfig( default_attributes={ "service_version": "1.2.3", "environment": "production", } ) plugin = BigQueryAgentAnalyticsPlugin( project_id, dataset_id, config=config ) ```
1 parent 1063fa5 commit 34fea3e

File tree

2 files changed

+140
-2
lines changed

2 files changed

+140
-2
lines changed

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -382,6 +382,10 @@ class BigQueryLoggerConfig:
382382
shutdown_timeout: Max time to wait for shutdown.
383383
queue_max_size: Max size of the in-memory queue.
384384
content_formatter: Optional custom formatter for content.
385+
gcs_bucket_name: GCS bucket name for offloading large content.
386+
connection_id: BigQuery connection ID for ObjectRef authorization.
387+
default_attributes: Static key-value pairs included in every event's
388+
attributes. Useful for service version, environment, etc.
385389
"""
386390

387391
enabled: bool = True
@@ -408,6 +412,10 @@ class BigQueryLoggerConfig:
408412
# If provided, this connection ID will be used as the authorizer for ObjectRef columns.
409413
# Format: "location.connection_id" (e.g. "us.my-connection")
410414
connection_id: Optional[str] = None
415+
# If provided, these key-value pairs will be merged into every event's attributes.
416+
# Useful for adding static metadata like service version, deployment environment, etc.
417+
# Event-specific attributes will override these if there are conflicts.
418+
default_attributes: Optional[dict[str, Any]] = None
411419

412420

413421
# ==============================================================================
@@ -1883,12 +1891,18 @@ async def _log_event(
18831891
# Fallback if it couldn't be converted to dict
18841892
kwargs["usage_metadata"] = usage_metadata
18851893

1894+
# Merge default_attributes first, then let event-specific kwargs override
1895+
if self.config.default_attributes:
1896+
merged_attributes = {**self.config.default_attributes, **kwargs}
1897+
else:
1898+
merged_attributes = kwargs
1899+
18861900
# Serialize remaining kwargs to JSON string for attributes
18871901
try:
1888-
attributes_json = json.dumps(kwargs)
1902+
attributes_json = json.dumps(merged_attributes)
18891903
except (TypeError, ValueError):
18901904
# Fallback for non-serializable objects
1891-
attributes_json = json.dumps(kwargs, default=str)
1905+
attributes_json = json.dumps(merged_attributes, default=str)
18921906

18931907
row = {
18941908
"timestamp": timestamp,

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2086,3 +2086,127 @@ async def test_flush_mechanism(
20862086
mock_write_client, dummy_arrow_schema
20872087
)
20882088
assert log_entry["event_type"] == "INVOCATION_STARTING"
2089+
2090+
@pytest.mark.asyncio
2091+
async def test_default_attributes_included_in_events(
2092+
self,
2093+
mock_write_client,
2094+
invocation_context,
2095+
mock_auth_default,
2096+
mock_bq_client,
2097+
mock_to_arrow_schema,
2098+
dummy_arrow_schema,
2099+
mock_asyncio_to_thread,
2100+
):
2101+
"""Test that default_attributes are included in every logged event."""
2102+
default_attrs = {
2103+
"service_version": "1.2.3",
2104+
"environment": "production",
2105+
"deployment_id": "deploy-abc",
2106+
}
2107+
config = BigQueryLoggerConfig(default_attributes=default_attrs)
2108+
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
2109+
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
2110+
)
2111+
await plugin._ensure_started()
2112+
mock_write_client.append_rows.reset_mock()
2113+
2114+
user_message = types.Content(parts=[types.Part(text="Hello")])
2115+
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
2116+
await plugin.on_user_message_callback(
2117+
invocation_context=invocation_context, user_message=user_message
2118+
)
2119+
await asyncio.sleep(0.01)
2120+
2121+
mock_write_client.append_rows.assert_called_once()
2122+
log_entry = await _get_captured_event_dict_async(
2123+
mock_write_client, dummy_arrow_schema
2124+
)
2125+
2126+
# Verify default attributes are in the attributes field
2127+
attributes = json.loads(log_entry["attributes"])
2128+
assert attributes["service_version"] == "1.2.3"
2129+
assert attributes["environment"] == "production"
2130+
assert attributes["deployment_id"] == "deploy-abc"
2131+
2132+
@pytest.mark.asyncio
2133+
async def test_default_attributes_overridden_by_event_attributes(
2134+
self,
2135+
mock_write_client,
2136+
callback_context,
2137+
mock_auth_default,
2138+
mock_bq_client,
2139+
mock_to_arrow_schema,
2140+
dummy_arrow_schema,
2141+
mock_asyncio_to_thread,
2142+
):
2143+
"""Test that event-specific attributes override default_attributes."""
2144+
default_attrs = {
2145+
"service_version": "1.2.3",
2146+
"model": "default-model",
2147+
}
2148+
config = BigQueryLoggerConfig(default_attributes=default_attrs)
2149+
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
2150+
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
2151+
)
2152+
await plugin._ensure_started()
2153+
mock_write_client.append_rows.reset_mock()
2154+
2155+
# LLM request will add its own "model" attribute which should override the default
2156+
llm_request = llm_request_lib.LlmRequest(
2157+
model="gemini-pro",
2158+
contents=[types.Content(parts=[types.Part(text="Hi")])],
2159+
)
2160+
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
2161+
await plugin.before_model_callback(
2162+
callback_context=callback_context, llm_request=llm_request
2163+
)
2164+
await asyncio.sleep(0.01)
2165+
2166+
mock_write_client.append_rows.assert_called_once()
2167+
log_entry = await _get_captured_event_dict_async(
2168+
mock_write_client, dummy_arrow_schema
2169+
)
2170+
2171+
attributes = json.loads(log_entry["attributes"])
2172+
# service_version should come from default_attributes
2173+
assert attributes["service_version"] == "1.2.3"
2174+
# model should be overridden by the event-specific value
2175+
assert attributes["model"] == "gemini-pro"
2176+
2177+
@pytest.mark.asyncio
2178+
async def test_default_attributes_none_does_not_affect_events(
2179+
self,
2180+
mock_write_client,
2181+
invocation_context,
2182+
mock_auth_default,
2183+
mock_bq_client,
2184+
mock_to_arrow_schema,
2185+
dummy_arrow_schema,
2186+
mock_asyncio_to_thread,
2187+
):
2188+
"""Test that when default_attributes is None, events work normally."""
2189+
config = BigQueryLoggerConfig(default_attributes=None)
2190+
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
2191+
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
2192+
)
2193+
await plugin._ensure_started()
2194+
mock_write_client.append_rows.reset_mock()
2195+
2196+
user_message = types.Content(parts=[types.Part(text="Hello")])
2197+
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
2198+
await plugin.on_user_message_callback(
2199+
invocation_context=invocation_context, user_message=user_message
2200+
)
2201+
await asyncio.sleep(0.01)
2202+
2203+
mock_write_client.append_rows.assert_called_once()
2204+
log_entry = await _get_captured_event_dict_async(
2205+
mock_write_client, dummy_arrow_schema
2206+
)
2207+
2208+
# Verify event was logged successfully with normal attributes
2209+
assert log_entry["event_type"] == "USER_MESSAGE_RECEIVED"
2210+
# Attributes should only contain root_agent_name (added by plugin)
2211+
attributes = json.loads(log_entry["attributes"])
2212+
assert "root_agent_name" in attributes

0 commit comments

Comments
 (0)