Skip to content

Commit 99be03e

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 21f63f6 commit 99be03e

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
@@ -350,6 +350,10 @@ class BigQueryLoggerConfig:
350350
shutdown_timeout: Max time to wait for shutdown.
351351
queue_max_size: Max size of the in-memory queue.
352352
content_formatter: Optional custom formatter for content.
353+
gcs_bucket_name: GCS bucket name for offloading large content.
354+
connection_id: BigQuery connection ID for ObjectRef authorization.
355+
default_attributes: Static key-value pairs included in every event's
356+
attributes. Useful for service version, environment, etc.
353357
"""
354358

355359
enabled: bool = True
@@ -376,6 +380,10 @@ class BigQueryLoggerConfig:
376380
# If provided, this connection ID will be used as the authorizer for ObjectRef columns.
377381
# Format: "location.connection_id" (e.g. "us.my-connection")
378382
connection_id: Optional[str] = None
383+
# If provided, these key-value pairs will be merged into every event's attributes.
384+
# Useful for adding static metadata like service version, deployment environment, etc.
385+
# Event-specific attributes will override these if there are conflicts.
386+
default_attributes: Optional[dict[str, Any]] = None
379387

380388

381389
# ==============================================================================
@@ -1713,12 +1721,18 @@ async def _log_event(
17131721
# Fallback if it couldn't be converted to dict
17141722
kwargs["usage_metadata"] = usage_metadata
17151723

1724+
# Merge default_attributes first, then let event-specific kwargs override
1725+
if self.config.default_attributes:
1726+
merged_attributes = {**self.config.default_attributes, **kwargs}
1727+
else:
1728+
merged_attributes = kwargs
1729+
17161730
# Serialize remaining kwargs to JSON string for attributes
17171731
try:
1718-
attributes_json = json.dumps(kwargs)
1732+
attributes_json = json.dumps(merged_attributes)
17191733
except (TypeError, ValueError):
17201734
# Fallback for non-serializable objects
1721-
attributes_json = json.dumps(kwargs, default=str)
1735+
attributes_json = json.dumps(merged_attributes, default=str)
17221736

17231737
row = {
17241738
"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
@@ -1953,3 +1953,127 @@ class LocalIncident:
19531953
content_json = json.loads(log_entry["content"])
19541954
assert content_json["result"]["id"] == "inc-123"
19551955
assert content_json["result"]["kpi_missed"][0]["kpi"] == "latency"
1956+
1957+
@pytest.mark.asyncio
1958+
async def test_default_attributes_included_in_events(
1959+
self,
1960+
mock_write_client,
1961+
invocation_context,
1962+
mock_auth_default,
1963+
mock_bq_client,
1964+
mock_to_arrow_schema,
1965+
dummy_arrow_schema,
1966+
mock_asyncio_to_thread,
1967+
):
1968+
"""Test that default_attributes are included in every logged event."""
1969+
default_attrs = {
1970+
"service_version": "1.2.3",
1971+
"environment": "production",
1972+
"deployment_id": "deploy-abc",
1973+
}
1974+
config = BigQueryLoggerConfig(default_attributes=default_attrs)
1975+
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
1976+
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
1977+
)
1978+
await plugin._ensure_started()
1979+
mock_write_client.append_rows.reset_mock()
1980+
1981+
user_message = types.Content(parts=[types.Part(text="Hello")])
1982+
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
1983+
await plugin.on_user_message_callback(
1984+
invocation_context=invocation_context, user_message=user_message
1985+
)
1986+
await asyncio.sleep(0.01)
1987+
1988+
mock_write_client.append_rows.assert_called_once()
1989+
log_entry = await _get_captured_event_dict_async(
1990+
mock_write_client, dummy_arrow_schema
1991+
)
1992+
1993+
# Verify default attributes are in the attributes field
1994+
attributes = json.loads(log_entry["attributes"])
1995+
assert attributes["service_version"] == "1.2.3"
1996+
assert attributes["environment"] == "production"
1997+
assert attributes["deployment_id"] == "deploy-abc"
1998+
1999+
@pytest.mark.asyncio
2000+
async def test_default_attributes_overridden_by_event_attributes(
2001+
self,
2002+
mock_write_client,
2003+
callback_context,
2004+
mock_auth_default,
2005+
mock_bq_client,
2006+
mock_to_arrow_schema,
2007+
dummy_arrow_schema,
2008+
mock_asyncio_to_thread,
2009+
):
2010+
"""Test that event-specific attributes override default_attributes."""
2011+
default_attrs = {
2012+
"service_version": "1.2.3",
2013+
"model": "default-model",
2014+
}
2015+
config = BigQueryLoggerConfig(default_attributes=default_attrs)
2016+
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
2017+
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
2018+
)
2019+
await plugin._ensure_started()
2020+
mock_write_client.append_rows.reset_mock()
2021+
2022+
# LLM request will add its own "model" attribute which should override the default
2023+
llm_request = llm_request_lib.LlmRequest(
2024+
model="gemini-pro",
2025+
contents=[types.Content(parts=[types.Part(text="Hi")])],
2026+
)
2027+
bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context)
2028+
await plugin.before_model_callback(
2029+
callback_context=callback_context, llm_request=llm_request
2030+
)
2031+
await asyncio.sleep(0.01)
2032+
2033+
mock_write_client.append_rows.assert_called_once()
2034+
log_entry = await _get_captured_event_dict_async(
2035+
mock_write_client, dummy_arrow_schema
2036+
)
2037+
2038+
attributes = json.loads(log_entry["attributes"])
2039+
# service_version should come from default_attributes
2040+
assert attributes["service_version"] == "1.2.3"
2041+
# model should be overridden by the event-specific value
2042+
assert attributes["model"] == "gemini-pro"
2043+
2044+
@pytest.mark.asyncio
2045+
async def test_default_attributes_none_does_not_affect_events(
2046+
self,
2047+
mock_write_client,
2048+
invocation_context,
2049+
mock_auth_default,
2050+
mock_bq_client,
2051+
mock_to_arrow_schema,
2052+
dummy_arrow_schema,
2053+
mock_asyncio_to_thread,
2054+
):
2055+
"""Test that when default_attributes is None, events work normally."""
2056+
config = BigQueryLoggerConfig(default_attributes=None)
2057+
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
2058+
PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config
2059+
)
2060+
await plugin._ensure_started()
2061+
mock_write_client.append_rows.reset_mock()
2062+
2063+
user_message = types.Content(parts=[types.Part(text="Hello")])
2064+
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
2065+
await plugin.on_user_message_callback(
2066+
invocation_context=invocation_context, user_message=user_message
2067+
)
2068+
await asyncio.sleep(0.01)
2069+
2070+
mock_write_client.append_rows.assert_called_once()
2071+
log_entry = await _get_captured_event_dict_async(
2072+
mock_write_client, dummy_arrow_schema
2073+
)
2074+
2075+
# Verify event was logged successfully with normal attributes
2076+
assert log_entry["event_type"] == "USER_MESSAGE_RECEIVED"
2077+
# Attributes should only contain root_agent_name (added by plugin)
2078+
attributes = json.loads(log_entry["attributes"])
2079+
assert "root_agent_name" in attributes

0 commit comments

Comments
 (0)