encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: Schema Enhancements with Descriptions, Partitioning, and Truncation Indicator

Browse Source 
This update enhances the BigQuery agent analytics plugin:

*   **Schema Field Descriptions:** The recommended BigQuery table schema now includes descriptions for each field, improving data understandability.
*   **Optimized Table Structure:** The plugin now creates the table with daily partitioning on `timestamp` and clustering on `event_type`, `agent`, and `user_id` by default.
*   **Truncation Flag:** A new boolean field `is_truncated` is added to the schema to show if the `content` was truncated.

PiperOrigin-RevId: 832436799
This commit is contained in:
Google Team Member
2025-11-14 13:28:45 -08:00
committed by Copybara-Service
parent 696852a280
commit 7c993b01d1
2 changed files with 298 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/plugins/bigquery_agent_analytics_plugin.py
+172 -39
View File
@@ -245,10 +245,19 @@ def _get_event_type(event: Event) -> str:
def _format_content(
content: Optional[types.Content], max_len: int = 500
) -> str:
"""Formats an Event content for logging."""
) -> tuple[str, bool]:
"""Formats an Event content for logging.
Args:
content: The Event content to format.
max_len: The maximum length of the text parts before truncation.
Returns:
A tuple containing the formatted content string and a boolean indicating if
the content was truncated.
"""
if not content or not content.parts:
return "None"
return "None", False
parts = []
for p in content.parts:
if p.text:
@@ -263,18 +272,33 @@ def _format_content(
parts.append(f"resp: {p.function_response.name}")
else:
parts.append("other")
return " | ".join(parts)
return " | ".join(parts), any(
len(p.text) > max_len for p in content.parts if p.text
)
def _format_args(args: dict[str, Any], max_len: int = 1000) -> str:
"""Formats tool arguments or results for logging."""
def _format_args(
args: dict[str, Any], *, max_len: int = 1000
) -> tuple[str, bool]:
"""Formats tool arguments or results for logging.
Args:
args: The tool arguments or results dictionary to format.
max_len: The maximum length of the output string before truncation.
Returns:
A tuple containing the JSON formatted string and a boolean indicating if
the content was truncated.
"""
if not args:
return "{}"
return "{}", False
try:
s = json.dumps(args)
except TypeError:
s = str(args)
return s[:max_len] + "..." if len(s) > max_len else s
if len(s) > max_len:
return s[:max_len] + "...", True
return s, False
class BigQueryAgentAnalyticsPlugin(BasePlugin):
@@ -322,29 +346,99 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
self._background_tasks: set[asyncio.Task] = set()
self._is_shutting_down = False
self._schema = [
bigquery.SchemaField("timestamp", "TIMESTAMP", mode="REQUIRED"),
bigquery.SchemaField("event_type", "STRING", mode="NULLABLE"),
bigquery.SchemaField("agent", "STRING", mode="NULLABLE"),
bigquery.SchemaField("session_id", "STRING", mode="NULLABLE"),
bigquery.SchemaField("invocation_id", "STRING", mode="NULLABLE"),
bigquery.SchemaField("user_id", "STRING", mode="NULLABLE"),
bigquery.SchemaField("content", "STRING", mode="NULLABLE"),
bigquery.SchemaField("error_message", "STRING", mode="NULLABLE"),
bigquery.SchemaField(
"timestamp",
"TIMESTAMP",
mode="REQUIRED",
description="The UTC time at which the event was logged.",
),
bigquery.SchemaField(
"event_type",
"STRING",
mode="NULLABLE",
description=(
"Indicates the type of event being logged (e.g., 'LLM_REQUEST',"
" 'TOOL_COMPLETED')."
),
),
bigquery.SchemaField(
"agent",
"STRING",
mode="NULLABLE",
description=(
"The name of the ADK agent or author associated with the event."
),
),
bigquery.SchemaField(
"session_id",
"STRING",
mode="NULLABLE",
description=(
"A unique identifier to group events within a single"
" conversation or user session."
),
),
bigquery.SchemaField(
"invocation_id",
"STRING",
mode="NULLABLE",
description=(
"A unique identifier for each individual agent execution or"
" turn within a session."
),
),
bigquery.SchemaField(
"user_id",
"STRING",
mode="NULLABLE",
description=(
"The identifier of the user associated with the current"
" session."
),
),
bigquery.SchemaField(
"content",
"STRING",
mode="NULLABLE",
description=(
"The event-specific data (payload). Format varies by"
" event_type."
),
),
bigquery.SchemaField(
"error_message",
"STRING",
mode="NULLABLE",
description=(
"Populated if an error occurs during the processing of the"
" event."
),
),
bigquery.SchemaField(
"is_truncated",
"BOOLEAN",
mode="NULLABLE",
description=(
"Indicates if the content field was truncated due to size"
" limits."
),
),
]
def _format_content_safely(
self, content: Optional[types.Content]
) -> str | None:
) -> tuple[str | None, bool]:
"""Formats content using self._config.content_formatter or _format_content, catching errors."""
if content is None:
return None
return None, False
try:
if self._config.content_formatter:
return self._config.content_formatter(content)
# Custom formatter: we assume no truncation or we can't know.
return self._config.content_formatter(content), False
return _format_content(content, max_len=self._config.max_content_length)
except Exception as e:
logging.warning(f"Content formatter failed: {e}")
return "[FORMATTING FAILED]"
return "[FORMATTING FAILED]", False
async def _ensure_init(self):
"""Ensures BigQuery clients are initialized."""
@@ -375,6 +469,10 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
f"{self._project_id}.{self._dataset_id}.{self._table_id}",
schema=self._schema,
)
table.time_partitioning = bigquery.TimePartitioning(
type_="DAY", field="timestamp"
)
table.clustering_fields = ["event_type", "agent", "user_id"]
self._bq_client.create_table(table, exists_ok=True)
logging.info(
"BQ Plugin: Dataset %s and Table %s ensured to exist.",
@@ -462,6 +560,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
"user_id": None,
"content": None,
"error_message": None,
"is_truncated": False,
}
row.update(data)
@@ -519,13 +618,15 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
user_message: types.Content,
) -> None:
"""Callback for user messages."""
content, truncated = self._format_content_safely(user_message)
await self._log({
"event_type": "USER_MESSAGE_RECEIVED",
"agent": invocation_context.agent.name,
"session_id": invocation_context.session.id,
"invocation_id": invocation_context.invocation_id,
"user_id": invocation_context.session.user_id,
"content": f"User Content: {self._format_content_safely(user_message)}",
"content": f"User Content: {content}",
"is_truncated": truncated,
})
async def before_run_callback(
@@ -544,15 +645,17 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
self, *, invocation_context: InvocationContext, event: Event
) -> None:
"""Callback for agent events."""
content, truncated = self._format_content_safely(event.content)
await self._log({
"event_type": _get_event_type(event),
"agent": event.author,
"session_id": invocation_context.session.id,
"invocation_id": invocation_context.invocation_id,
"user_id": invocation_context.session.user_id,
"content": self._format_content_safely(event.content),
"content": content,
"error_message": event.error_message,
"timestamp": datetime.fromtimestamp(event.timestamp, timezone.utc),
"is_truncated": truncated,
})
async def after_run_callback(
@@ -600,10 +703,15 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
content_parts = [
f"Model: {llm_request.model or 'default'}",
]
is_truncated = False
if contents := getattr(llm_request, "contents", None):
prompt_str = " | ".join(
[f"{c.role}: {self._format_content_safely(c)}" for c in contents]
)
prompt_parts = []
for c in contents:
c_str, c_trunc = self._format_content_safely(c)
prompt_parts.append(f"{c.role}: {c_str}")
if c_trunc:
is_truncated = True
prompt_str = " | ".join(prompt_parts)
content_parts.append(f"Prompt: {prompt_str}")
system_instruction_text = "None"
if llm_request.config and llm_request.config.system_instruction:
@@ -656,6 +764,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
max_len = self._config.max_content_length
if len(final_content) > max_len:
final_content = final_content[:max_len] + "..."
is_truncated = True
await self._log({
"event_type": "LLM_REQUEST",
"agent": callback_context.agent_name,
@@ -663,6 +772,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
"invocation_id": callback_context.invocation_id,
"user_id": callback_context.session.user_id,
"content": final_content,
"is_truncated": is_truncated,
})
async def after_model_callback(
@@ -672,6 +782,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
content_parts = []
content = llm_response.content
is_tool_call = False
is_truncated = False
if content and content.parts:
is_tool_call = any(part.function_call for part in content.parts)
@@ -685,8 +796,12 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
]
content_parts.append(f"Tool Name: {', '.join(fc_names)}")
else:
text_content = self._format_content_safely(llm_response.content)
text_content, truncated = self._format_content_safely(
llm_response.content
)
content_parts.append(f"Tool Name: text_response, {text_content}")
if truncated:
is_truncated = True
if llm_response.usage_metadata:
prompt_tokens = getattr(
@@ -713,6 +828,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
"user_id": callback_context.session.user_id,
"content": final_content,
"error_message": llm_response.error_message,
"is_truncated": is_truncated,
})
async def before_tool_callback(
@@ -723,17 +839,24 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
tool_context: ToolContext,
) -> None:
"""Callback before tool call."""
args_str, truncated = _format_args(
tool_args, max_len=self._config.max_content_length
)
content = (
f"Tool Name: {tool.name}, Description: {tool.description},"
f" Arguments: {args_str}"
)
if len(content) > self._config.max_content_length:
content = content[: self._config.max_content_length] + "..."
truncated = True
await self._log({
"event_type": "TOOL_STARTING",
"agent": tool_context.agent_name,
"session_id": tool_context.session.id,
"invocation_id": tool_context.invocation_id,
"user_id": tool_context.session.user_id,
"content": (
f"Tool Name: {tool.name}, Description: {tool.description},"
" Arguments:"
f" {_format_args(tool_args, max_len=self._config.max_content_length)}"
),
"content": content,
"is_truncated": truncated,
})
async def after_tool_callback(
@@ -745,16 +868,21 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
result: dict[str, Any],
) -> None:
"""Callback after tool call."""
result_str, truncated = _format_args(
result, max_len=self._config.max_content_length
)
content = f"Tool Name: {tool.name}, Result: {result_str}"
if len(content) > self._config.max_content_length:
content = content[: self._config.max_content_length] + "..."
truncated = True
await self._log({
"event_type": "TOOL_COMPLETED",
"agent": tool_context.agent_name,
"session_id": tool_context.session.id,
"invocation_id": tool_context.invocation_id,
"user_id": tool_context.session.user_id,
"content": (
f"Tool Name: {tool.name}, Result:"
f" {_format_args(result, max_len=self._config.max_content_length)}"
),
"content": content,
"is_truncated": truncated,
})
async def on_model_error_callback(
@@ -783,15 +911,20 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin):
error: Exception,
) -> None:
"""Callback for tool errors."""
args_str, truncated = _format_args(
tool_args, max_len=self._config.max_content_length
)
content = f"Tool Name: {tool.name}, Arguments: {args_str}"
if len(content) > self._config.max_content_length:
content = content[: self._config.max_content_length] + "..."
truncated = True
await self._log({
"event_type": "TOOL_ERROR",
"agent": tool_context.agent_name,
"session_id": tool_context.session.id,
"invocation_id": tool_context.invocation_id,
"user_id": tool_context.session.user_id,
"content": (
f"Tool Name: {tool.name}, Arguments:"
f" {_format_args(tool_args, max_len=self._config.max_content_length)}"
),
"content": content,
"error_message": str(error),
"is_truncated": truncated,
})
tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py
+126 -7
View File
@@ -154,6 +154,7 @@ def dummy_arrow_schema():
pa.field("user_id", pa.string(), nullable=True),
pa.field("content", pa.string(), nullable=True),
pa.field("error_message", pa.string(), nullable=True),
pa.field("is_truncated", pa.bool_(), nullable=True),
])
@@ -233,6 +234,7 @@ def _assert_common_fields(log_entry, event_type, agent="MyTestAgent"):
assert log_entry["user_id"] == "user-456"
assert "timestamp" in log_entry
assert isinstance(log_entry["timestamp"], datetime.datetime)
assert "is_truncated" in log_entry
# --- Test Class ---
@@ -424,6 +426,7 @@ class TestBigQueryAgentAnalyticsPlugin:
log_entry["content"]
== "User Content: text: '1234567890123456789012345678901234567890...' "
)
assert log_entry["is_truncated"]
mock_write_client.append_rows.reset_mock()
# Test before_model_callback full content truncation
@@ -448,6 +451,7 @@ class TestBigQueryAgentAnalyticsPlugin:
# Truncated to 40 chars + ...:
expected_content = "Model: gemini-pro | Prompt: user: text: ..."
assert log_entry["content"] == expected_content
assert log_entry["is_truncated"]
@pytest.mark.asyncio
async def test_max_content_length_tool_args(
@@ -460,7 +464,7 @@ class TestBigQueryAgentAnalyticsPlugin:
dummy_arrow_schema,
mock_asyncio_to_thread,
):
config = BigQueryLoggerConfig(max_content_length=10)
config = BigQueryLoggerConfig(max_content_length=80)
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
PROJECT_ID, DATASET_ID, TABLE_ID, config
)
@@ -473,19 +477,104 @@ class TestBigQueryAgentAnalyticsPlugin:
type(mock_tool).name = mock.PropertyMock(return_value="MyTool")
type(mock_tool).description = mock.PropertyMock(return_value="Description")
# Args length > 10
# {"param": "long_value"} is ~24 chars
# Args length > 80
# {"param": "A" * 50} is ~60 chars.
# Prefix is ~57 chars. Total ~117 chars.
await plugin.before_tool_callback(
tool=mock_tool,
tool_args={"param": "long_value"},
tool_args={"param": "A" * 50},
tool_context=tool_context,
)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
# JSON string: '{"param": "long_value"}'
# Truncated to 10: '{"param": ...'
assert 'Arguments: {"param": ...' in log_entry["content"]
assert 'Arguments: {"param": "AAAAA' in log_entry["content"]
assert log_entry["content"].endswith("...")
assert len(log_entry["content"]) == 83 # 80 + 3 dots
assert log_entry["is_truncated"]
@pytest.mark.asyncio
async def test_max_content_length_tool_result(
self,
mock_write_client,
tool_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
config = BigQueryLoggerConfig(max_content_length=80)
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
PROJECT_ID, DATASET_ID, TABLE_ID, config
)
await plugin._ensure_init()
mock_write_client.append_rows.reset_mock()
mock_tool = mock.create_autospec(
base_tool_lib.BaseTool, instance=True, spec_set=True
)
type(mock_tool).name = mock.PropertyMock(return_value="MyTool")
# Result length > 80
# {"res": "A" * 60} is ~70 chars.
# Prefix is ~27 chars. Total ~97 chars.
await plugin.after_tool_callback(
tool=mock_tool,
tool_args={},
tool_context=tool_context,
result={"res": "A" * 60},
)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
assert 'Result: {"res": "AAAAA' in log_entry["content"]
assert log_entry["content"].endswith("...")
assert len(log_entry["content"]) == 83 # 80 + 3 dots
assert log_entry["is_truncated"]
@pytest.mark.asyncio
async def test_max_content_length_tool_error(
self,
mock_write_client,
tool_context,
mock_auth_default,
mock_bq_client,
mock_to_arrow_schema,
dummy_arrow_schema,
mock_asyncio_to_thread,
):
config = BigQueryLoggerConfig(max_content_length=80)
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
PROJECT_ID, DATASET_ID, TABLE_ID, config
)
await plugin._ensure_init()
mock_write_client.append_rows.reset_mock()
mock_tool = mock.create_autospec(
base_tool_lib.BaseTool, instance=True, spec_set=True
)
type(mock_tool).name = mock.PropertyMock(return_value="MyTool")
# Args length > 80
# {"arg": "A" * 60} is ~70 chars.
# Prefix is ~28 chars. Total ~98 chars.
await plugin.on_tool_error_callback(
tool=mock_tool,
tool_args={"arg": "A" * 60},
tool_context=tool_context,
error=ValueError("Oops"),
)
await asyncio.sleep(0.01)
mock_write_client.append_rows.assert_called_once()
log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
assert 'Arguments: {"arg": "AAAAA' in log_entry["content"]
assert log_entry["content"].endswith("...")
assert len(log_entry["content"]) == 83 # 80 + 3 dots
assert log_entry["is_truncated"]
@pytest.mark.asyncio
async def test_on_user_message_callback_logs_correctly(
@@ -503,6 +592,7 @@ class TestBigQueryAgentAnalyticsPlugin:
log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema)
_assert_common_fields(log_entry, "USER_MESSAGE_RECEIVED")
assert log_entry["content"] == "User Content: text: 'What is up?'"
assert not log_entry["is_truncated"]
@pytest.mark.asyncio
async def test_on_event_callback_tool_call(
@@ -857,3 +947,32 @@ class TestBigQueryAgentAnalyticsPlugin:
== 'Tool Name: MyTool, Arguments: {"param": "value"}'
)
assert log_entry["error_message"] == "Tool timed out"
@pytest.mark.asyncio
async def test_table_creation_options(
self,
mock_auth_default,
mock_bq_client,
mock_write_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
PROJECT_ID, DATASET_ID, TABLE_ID
)
await plugin._ensure_init()
# Verify create_table was called with correct table options
mock_bq_client.create_table.assert_called_once()
call_args = mock_bq_client.create_table.call_args
table_arg = call_args[0][0]
assert isinstance(table_arg, bigquery.Table)
assert table_arg.time_partitioning.type_ == "DAY"
assert table_arg.time_partitioning.field == "timestamp"
assert table_arg.clustering_fields == ["event_type", "agent", "user_id"]
# Verify schema descriptions are present (spot check)
timestamp_field = next(f for f in table_arg.schema if f.name == "timestamp")
assert (
timestamp_field.description
== "The UTC time at which the event was logged."
)