mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
156d235479
commit
0b56f2287b
@@ -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"
|
||||||
|
|||||||
Reference in New Issue
Block a user