From b7dbfed4a3d4a0165e2c6e51594d1f547bec89d3 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Mon, 27 Oct 2025 15:30:15 -0700 Subject: [PATCH] feat: Add BigQueryLoggingPlugin for event logging to BigQuery Introduces the `BigQueryLoggingPlugin` for capturing and sending ADK lifecycle events to Google BigQuery. This allows for persistent storage and analysis of agent and tool interactions. The plugin supports asynchronous logging, automatic dataset/table creation, and comprehensive event capture. Also refactors common formatting utilities (_format_content, _format_args) for shared use. PiperOrigin-RevId: 824703739 --- .../adk/plugins/bigquery_logging_plugin.py | 550 ++++++++++++++++++ .../plugins/test_bigquery_logging_plugin.py | 403 +++++++++++++ 2 files changed, 953 insertions(+) create mode 100644 src/google/adk/plugins/bigquery_logging_plugin.py create mode 100644 tests/unittests/plugins/test_bigquery_logging_plugin.py diff --git a/src/google/adk/plugins/bigquery_logging_plugin.py b/src/google/adk/plugins/bigquery_logging_plugin.py new file mode 100644 index 00000000..9740a745 --- /dev/null +++ b/src/google/adk/plugins/bigquery_logging_plugin.py @@ -0,0 +1,550 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import asyncio +from datetime import datetime +from datetime import timezone +import json +import logging +import threading +from typing import Any +from typing import Optional +from typing import TYPE_CHECKING + +import google.api_core.client_info +import google.auth +from google.auth import exceptions as auth_exceptions +from google.cloud import bigquery +from google.cloud import exceptions as cloud_exceptions +from google.genai import types + +from .. import version +from ..agents.base_agent import BaseAgent +from ..agents.callback_context import CallbackContext +from ..events.event import Event +from ..models.llm_request import LlmRequest +from ..models.llm_response import LlmResponse +from ..tools.base_tool import BaseTool +from ..tools.tool_context import ToolContext +from .base_plugin import BasePlugin + +if TYPE_CHECKING: + from ..agents.invocation_context import InvocationContext + + +def _get_event_type(event: Event) -> str: + if event.author == "user": + return "USER_INPUT" + if event.get_function_calls(): + return "TOOL_CALL" + if event.get_function_responses(): + return "TOOL_RESULT" + if event.content and event.content.parts: + return "MODEL_RESPONSE" + if event.error_message: + return "ERROR" + return "SYSTEM" # Fallback for other event types + + +def _format_content( + content: Optional[types.Content], max_length: int = 200 +) -> str: + """Format content for logging, truncating if too long.""" + if not content or not content.parts: + return "None" + parts = [] + for part in content.parts: + if part.text: + text = part.text.strip() + if len(text) > max_length: + text = text[:max_length] + "..." + parts.append(f"text: '{text}'") + elif part.function_call: + parts.append(f"function_call: {part.function_call.name}") + elif part.function_response: + parts.append(f"function_response: {part.function_response.name}") + elif part.code_execution_result: + parts.append("code_execution_result") + else: + parts.append("other_part") + return " | ".join(parts) + + +def _format_args(args: dict[str, Any], max_length: int = 300) -> str: + """Format arguments dictionary for logging.""" + if not args: + return "{}" + formatted = str(args) + if len(formatted) > max_length: + formatted = formatted[:max_length] + "...}" + return formatted + + +class BigQueryAgentAnalyticsPlugin(BasePlugin): + """A plugin that logs ADK events to a BigQuery table. + + This plugin captures critical events during an agent invocation and logs them + as structured data to the specified BigQuery table. This allows for + persistent storage, auditing, and analysis of agent interactions. + + The plugin logs the following information at each callback point: + - User messages and invocation context + - Agent execution flow (start and completion) + - LLM requests and responses (including token usage in content) + - Tool calls with arguments and results + - Events yielded by agents + - Errors during model and tool execution + + Each log entry includes a timestamp, event type, agent name, session ID, + invocation ID, user ID, content payload, and any error messages. + """ + + def __init__( + self, + project_id: str, + dataset_id: str = "adk_agent_logs", + table_id: str = "agent_events", + **kwargs, + ): + super().__init__(name=kwargs.get("name", "BigQueryAgentAnalyticsPlugin")) + self._project_id = project_id + self._dataset_id = dataset_id + self._table_id = table_id + self._bq_client: bigquery.Client | None = None + self._client_init_lock = threading.Lock() + self._init_done = False + self._init_succeeded = False + logging.debug( + "DEBUG: BigQueryAgentAnalyticsPlugin INSTANTIATED (Name: %s)", self.name + ) + + def _ensure_initialized_sync(self): + """Synchronous initialization of BQ client and table.""" + with self._client_init_lock: + if self._init_done: + return + self._init_done = True + try: + credentials, _ = google.auth.default( + scopes=["https://www.googleapis.com/auth/bigquery"] + ) + client_info = google.api_core.client_info.ClientInfo( + user_agent=f"google-adk-plugin/{version.__version__}" + ) + self._bq_client = bigquery.Client( + project=self._project_id, + credentials=credentials, + client_info=client_info, + ) + logging.info( + "BigQuery client initialized for project %s", self._project_id + ) + dataset_ref = self._bq_client.dataset(self._dataset_id) + self._bq_client.create_dataset(dataset_ref, exists_ok=True) + logging.info("Dataset %s ensured to exist.", self._dataset_id) + table_ref = dataset_ref.table(self._table_id) + # Schema without separate token columns + schema = [ + bigquery.SchemaField("dataset_id", "STRING"), + bigquery.SchemaField("timestamp", "TIMESTAMP"), + bigquery.SchemaField("event_type", "STRING"), + bigquery.SchemaField("agent", "STRING"), + bigquery.SchemaField("session_id", "STRING"), + bigquery.SchemaField("invocation_id", "STRING"), + bigquery.SchemaField("user_id", "STRING"), + bigquery.SchemaField("content", "STRING"), + bigquery.SchemaField("error_message", "STRING"), + ] + table = bigquery.Table(table_ref, schema=schema) + self._bq_client.create_table(table, exists_ok=True) + logging.info("Table %s ensured to exist.", self._table_id) + self._init_succeeded = True + except ( + auth_exceptions.GoogleAuthError, + cloud_exceptions.GoogleCloudError, + ) as e: + logging.exception( + "Failed to initialize BigQuery client or table: %s", e + ) + self._init_succeeded = False + + async def _log_to_bigquery_async(self, event_dict: dict[str, Any]): + def _sync_log(): + self._ensure_initialized_sync() + if not self._init_succeeded or not self._bq_client: + return + table_ref = self._bq_client.dataset(self._dataset_id).table( + self._table_id + ) + default_row = { + "dataset_id": None, + "timestamp": datetime.now(timezone.utc).isoformat(), + "event_type": None, + "agent": None, + "session_id": None, + "invocation_id": None, + "user_id": None, + "content": None, + "error_message": None, + } + insert_row = {**default_row, **event_dict} + + errors = self._bq_client.insert_rows_json(table_ref, [insert_row]) + if errors: + logging.error( + "Errors occurred while inserting to BigQuery table %s.%s: %s", + self._dataset_id, + self._table_id, + errors, + ) + + try: + await asyncio.to_thread(_sync_log) + except ( + cloud_exceptions.GoogleCloudError, + auth_exceptions.GoogleAuthError, + ) as e: + logging.exception("Failed to log to BigQuery: %s", e) + + async def on_user_message_callback( + self, + *, + invocation_context: InvocationContext, + user_message: types.Content, + ) -> Optional[types.Content]: + """Log user message and invocation start.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "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: {_format_content(user_message)}", + } + await self._log_to_bigquery_async(event_dict) + return None + + async def before_run_callback( + self, *, invocation_context: InvocationContext + ) -> Optional[types.Content]: + """Log invocation start.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "event_type": "INVOCATION_STARTING", + "agent": invocation_context.agent.name, + "session_id": invocation_context.session.id, + "invocation_id": invocation_context.invocation_id, + "user_id": invocation_context.session.user_id, + } + await self._log_to_bigquery_async(event_dict) + return None + + async def on_event_callback( + self, *, invocation_context: InvocationContext, event: Event + ) -> Optional[Event]: + """Logs event data to BigQuery.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.fromtimestamp( + event.timestamp, timezone.utc + ).isoformat(), + "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": ( + json.dumps( + [part.model_dump(mode="json") for part in event.content.parts] + ) + if event.content and event.content.parts + else None + ), + "error_message": event.error_message, + } + await self._log_to_bigquery_async(event_dict) + return None + + async def after_run_callback( + self, *, invocation_context: InvocationContext + ) -> Optional[None]: + """Log invocation completion.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "event_type": "INVOCATION_COMPLETED", + "agent": invocation_context.agent.name, + "session_id": invocation_context.session.id, + "invocation_id": invocation_context.invocation_id, + "user_id": invocation_context.session.user_id, + } + await self._log_to_bigquery_async(event_dict) + return None + + async def before_agent_callback( + self, *, agent: BaseAgent, callback_context: CallbackContext + ) -> Optional[types.Content]: + """Log agent execution start.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "event_type": "AGENT_STARTING", + "agent": agent.name, + "session_id": callback_context.session.id, + "invocation_id": callback_context.invocation_id, + "user_id": callback_context.session.user_id, + "content": f"Agent Name: {callback_context.agent_name}", + } + await self._log_to_bigquery_async(event_dict) + return None + + async def after_agent_callback( + self, *, agent: BaseAgent, callback_context: CallbackContext + ) -> Optional[types.Content]: + """Log agent execution completion.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "event_type": "AGENT_COMPLETED", + "agent": agent.name, + "session_id": callback_context.session.id, + "invocation_id": callback_context.invocation_id, + "user_id": callback_context.session.user_id, + "content": f"Agent Name: {callback_context.agent_name}", + } + await self._log_to_bigquery_async(event_dict) + return None + + async def before_model_callback( + self, *, callback_context: CallbackContext, llm_request: LlmRequest + ) -> Optional[LlmResponse]: + """Log LLM request before sending to model, including the full system instruction.""" + + content_parts = [ + f"Model: {llm_request.model or 'default'}", + ] + + # Log Full System Instruction + system_instruction_text = "None" + if llm_request.config and hasattr(llm_request.config, "system_instruction"): + si = llm_request.config.system_instruction + if si: + if isinstance(si, str): + system_instruction_text = si + elif hasattr(si, "__iter__"): # Handles list, tuple, etc. of parts + # Join parts together to form the complete system instruction + system_instruction_text = "".join( + part.text for part in si if hasattr(part, "text") + ) + else: + system_instruction_text = str(si) + else: + system_instruction_text = "Empty" + + content_parts.append(f"System Prompt: {system_instruction_text}") + + # Log Generation Config Parameters + if llm_request.config: + config = llm_request.config + params_to_log = {} + if hasattr(config, "temperature") and config.temperature is not None: + params_to_log["temperature"] = config.temperature + if hasattr(config, "top_p") and config.top_p is not None: + params_to_log["top_p"] = config.top_p + if hasattr(config, "top_k") and config.top_k is not None: + params_to_log["top_k"] = config.top_k + if ( + hasattr(config, "max_output_tokens") + and config.max_output_tokens is not None + ): + params_to_log["max_output_tokens"] = config.max_output_tokens + + if params_to_log: + params_str = ", ".join([f"{k}={v}" for k, v in params_to_log.items()]) + content_parts.append(f"Params: {{{params_str}}}") + + if llm_request.tools_dict: + content_parts.append( + f"Available Tools: {list(llm_request.tools_dict.keys())}" + ) + + final_content = " | ".join(content_parts) + + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "event_type": "LLM_REQUEST", + "agent": callback_context.agent_name, + "session_id": callback_context.session.id, + "invocation_id": callback_context.invocation_id, + "user_id": callback_context.session.user_id, + "content": final_content, + } + await self._log_to_bigquery_async(event_dict) + return None + + async def after_model_callback( + self, *, callback_context: CallbackContext, llm_response: LlmResponse + ) -> Optional[LlmResponse]: + """Log LLM response after receiving from model.""" + content_parts = [] + content = llm_response.content + is_tool_call = False + if content and content.parts: + is_tool_call = any(part.function_call for part in content.parts) + + if is_tool_call: + # Explicitly state Tool Name + fc_names = [] + if content and content.parts: + fc_names = [ + part.function_call.name + for part in content.parts + if part.function_call + ] + content_parts.append(f"Tool Name: {', '.join(fc_names)}") + else: + # This is a text response + text_content = _format_content( + llm_response.content + ) # This returns something like "text: 'The actual message...'" + content_parts.append(f"Tool Name: text_response, {text_content}") + + if llm_response.usage_metadata: + prompt_tokens = getattr( + llm_response.usage_metadata, "prompt_token_count", "N/A" + ) + candidates_tokens = getattr( + llm_response.usage_metadata, "candidates_token_count", "N/A" + ) + total_tokens = getattr( + llm_response.usage_metadata, "total_token_count", "N/A" + ) + token_usage_str = ( + f"Token Usage: {{prompt: {prompt_tokens}, candidates:" + f" {candidates_tokens}, total: {total_tokens}}}" + ) + content_parts.append(token_usage_str) + + final_content = " | ".join(content_parts) + + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "event_type": "LLM_RESPONSE", + "agent": callback_context.agent_name, + "session_id": callback_context.session.id, + "invocation_id": callback_context.invocation_id, + "user_id": callback_context.session.user_id, + "content": final_content, + "error_message": ( + llm_response.error_message if llm_response.error_code else None + ), + } + await self._log_to_bigquery_async(event_dict) + return None + + async def before_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + ) -> Optional[None]: + """Log tool execution start.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "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}," + f" Arguments: {_format_args(tool_args)}" + ), + } + await self._log_to_bigquery_async(event_dict) + return None + + async def after_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + result: dict[str, Any], + ) -> None: + """Log tool execution completion.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "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: {_format_args(result)}", + } + await self._log_to_bigquery_async(event_dict) + return None + + async def on_model_error_callback( + self, + *, + callback_context: CallbackContext, + llm_request: LlmRequest, + error: Exception, + ) -> Optional[LlmResponse]: + """Log LLM error.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "event_type": "LLM_ERROR", + "agent": callback_context.agent_name, + "session_id": callback_context.session.id, + "invocation_id": callback_context.invocation_id, + "user_id": callback_context.session.user_id, + "error_message": str(error), + } + await self._log_to_bigquery_async(event_dict) + return None + + async def on_tool_error_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + error: Exception, + ) -> None: + """Log tool error.""" + event_dict = { + "dataset_id": self._dataset_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "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}", + "error_message": str(error), + } + await self._log_to_bigquery_async(event_dict) + return None diff --git a/tests/unittests/plugins/test_bigquery_logging_plugin.py b/tests/unittests/plugins/test_bigquery_logging_plugin.py new file mode 100644 index 00000000..557eb176 --- /dev/null +++ b/tests/unittests/plugins/test_bigquery_logging_plugin.py @@ -0,0 +1,403 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import datetime +import json +import logging +from unittest import mock + +from google.adk.agents import base_agent +from google.adk.agents import callback_context as callback_context_lib +from google.adk.agents import invocation_context as invocation_context_lib +from google.adk.events import event as event_lib +from google.adk.models import llm_request as llm_request_lib +from google.adk.models import llm_response as llm_response_lib +from google.adk.plugins import bigquery_logging_plugin +from google.adk.plugins import plugin_manager as plugin_manager_lib +from google.adk.sessions import base_session_service as base_session_service_lib +from google.adk.sessions import session as session_lib +from google.adk.tools import base_tool as base_tool_lib +from google.adk.tools import tool_context as tool_context_lib +import google.auth +from google.auth import exceptions as auth_exceptions +from google.cloud import bigquery +from google.genai import types +import pytest + + +class PluginTestBase: + """Base class for plugin tests with common context setup.""" + + def setup_method(self, method): + self.mock_session = mock.create_autospec(session_lib.Session, instance=True) + self.mock_session.id = "session-123" + self.mock_session.user_id = "user-456" + self.mock_session.app_name = "test_app" + self.mock_session.state = {} + self.mock_agent = mock.create_autospec(base_agent.BaseAgent, instance=True) + self.mock_agent.name = "MyTestAgent" + mock_session_service = mock.create_autospec( + base_session_service_lib.BaseSessionService, instance=True + ) + mock_plugin_manager = mock.create_autospec( + plugin_manager_lib.PluginManager, instance=True + ) + self.invocation_context = invocation_context_lib.InvocationContext( + agent=self.mock_agent, + session=self.mock_session, + invocation_id="inv-789", + session_service=mock_session_service, + plugin_manager=mock_plugin_manager, + ) + self.callback_context = callback_context_lib.CallbackContext( + invocation_context=self.invocation_context + ) + self.tool_context = tool_context_lib.ToolContext( + invocation_context=self.invocation_context + ) + + def teardown_method(self, method): + mock.patch.stopall() + + +class TestBigQueryAgentAnalyticsPlugin(PluginTestBase): + """Tests for the BigQueryAgentAnalyticsPlugin.""" + + def setup_method(self, method): + super().setup_method(method) + self.project_id = "test-gcp-project" + self.dataset_id = "adk_logs" + self.table_id = "agent_events" + + # Mock Google Auth default credentials + self._auth_patch = mock.patch.object(google.auth, "default", autospec=True) + self.mock_auth_default = self._auth_patch.start() + self.mock_auth_default.return_value = (mock.Mock(), self.project_id) + + # Mock BigQuery Client class + self._bq_client_patch = mock.patch.object(bigquery, "Client", autospec=True) + self.mock_bq_client_cls = self._bq_client_patch.start() + self.mock_bq_client = self.mock_bq_client_cls.return_value + self.mock_bq_client.create_dataset.return_value = None + self.mock_bq_client.create_table.return_value = None + self.mock_bq_client.insert_rows_json.return_value = [] # No errors + self.mock_table_ref = mock.Mock() + self.mock_table_ref.dataset_id = self.dataset_id + self.mock_table_ref.table_id = self.table_id + self.mock_dataset_ref = mock.Mock() + self.mock_dataset_ref.table.return_value = self.mock_table_ref + self.mock_bq_client.dataset.return_value = self.mock_dataset_ref + + # Patch asyncio.to_thread to run the function synchronously + self._asyncio_to_thread_patch = mock.patch( + "asyncio.to_thread", + side_effect=lambda func, *args, **kwargs: func(*args, **kwargs), + ) + self._asyncio_to_thread_patch.start() + + self.plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin( + project_id=self.project_id, + dataset_id=self.dataset_id, + table_id=self.table_id, + ) + # Trigger lazy initialization by calling an async method once. + asyncio.run(self.plugin._log_to_bigquery_async({"event_type": "INIT"})) + self.mock_bq_client.insert_rows_json.reset_mock() + + def _get_logged_entry(self): + """Helper to get the single logged entry from the mocked client.""" + self.mock_bq_client.insert_rows_json.assert_called_once() + args, _ = self.mock_bq_client.insert_rows_json.call_args + rows = args[1] + assert len(rows) == 1 + return rows[0] + + 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["agent"] == "MyTestAgent" + assert log_entry["session_id"] == "session-123" + assert log_entry["invocation_id"] == "inv-789" + assert log_entry["user_id"] == "user-456" + assert log_entry["timestamp"] is not None + + @pytest.mark.asyncio + async def test_on_user_message_callback_logs_correctly(self): + user_message = types.Content(parts=[types.Part(text="What is up?")]) + await self.plugin.on_user_message_callback( + invocation_context=self.invocation_context, user_message=user_message + ) + + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "USER_MESSAGE_RECEIVED") + assert log_entry["content"] == "User Content: text: 'What is up?'" + + @pytest.mark.asyncio + async def test_on_event_callback_tool_call(self): + tool_fc = types.FunctionCall(name="get_weather", args={"location": "Paris"}) + event = event_lib.Event( + author="MyTestAgent", + content=types.Content(parts=[types.Part(function_call=tool_fc)]), + timestamp=datetime.datetime( + 2025, 10, 22, 10, 0, 0, tzinfo=datetime.timezone.utc + ).timestamp(), + ) + await self.plugin.on_event_callback( + invocation_context=self.invocation_context, event=event + ) + + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "TOOL_CALL") + logged_content = json.loads(log_entry["content"]) + assert logged_content[0]["function_call"]["args"] == {"location": "Paris"} + assert logged_content[0]["function_call"]["name"] == "get_weather" + assert log_entry["timestamp"] == "2025-10-22T10:00:00+00:00" + + @pytest.mark.asyncio + async def test_on_event_callback_model_response(self): + event = event_lib.Event( + author="MyTestAgent", + content=types.Content(parts=[types.Part(text="Hello there!")]), + timestamp=datetime.datetime( + 2025, 10, 22, 11, 0, 0, tzinfo=datetime.timezone.utc + ).timestamp(), + ) + await self.plugin.on_event_callback( + invocation_context=self.invocation_context, event=event + ) + + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "MODEL_RESPONSE") + logged_content = json.loads(log_entry["content"]) + assert logged_content[0]["text"] == "Hello there!" + assert log_entry["timestamp"] == "2025-10-22T11:00:00+00:00" + + @pytest.mark.asyncio + async def test_bigquery_client_initialization_failure(self): + # Simulate auth failure + self.mock_auth_default.side_effect = auth_exceptions.GoogleAuthError( + "Auth failed" + ) + self.mock_bq_client.insert_rows_json.reset_mock() + + # Re-instantiate the plugin so init is re-attempted + plugin_with_fail = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin( + project_id=self.project_id, + dataset_id=self.dataset_id, + table_id=self.table_id, + ) + + # Trigger a callback; initialization happens lazily + with mock.patch.object(logging, "exception") as mock_log_exception: + await plugin_with_fail.before_run_callback( + invocation_context=self.invocation_context + ) + mock_log_exception.assert_called_once() + + # Ensure insert_rows_json was never called because init failed + self.mock_bq_client.insert_rows_json.assert_not_called() + + @pytest.mark.asyncio + async def test_bigquery_insert_error_does_not_raise(self): + # Simulate an insert error in the future result + self.mock_bq_client.insert_rows_json.return_value = [{"errors": ["error"]}] + + with mock.patch.object(logging, "error") as mock_log_error: + await self.plugin.on_user_message_callback( + invocation_context=self.invocation_context, + user_message=types.Content(parts=[types.Part(text="Test")]), + ) + # The plugin should handle the error internally without raising + mock_log_error.assert_called_with( + "Errors occurred while inserting to BigQuery table %s.%s: %s", + self.dataset_id, + self.table_id, + [{"errors": ["error"]}], + ) + + self.mock_bq_client.insert_rows_json.assert_called_once() + + @pytest.mark.asyncio + async def test_before_run_callback_logs_correctly(self): + await self.plugin.before_run_callback( + invocation_context=self.invocation_context + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "INVOCATION_STARTING") + assert log_entry["content"] is None + + @pytest.mark.asyncio + async def test_after_run_callback_logs_correctly(self): + await self.plugin.after_run_callback( + invocation_context=self.invocation_context + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "INVOCATION_COMPLETED") + assert log_entry["content"] is None + + @pytest.mark.asyncio + async def test_before_agent_callback_logs_correctly(self): + await self.plugin.before_agent_callback( + agent=self.mock_agent, callback_context=self.callback_context + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "AGENT_STARTING") + assert log_entry["content"] == "Agent Name: MyTestAgent" + + @pytest.mark.asyncio + async def test_after_agent_callback_logs_correctly(self): + await self.plugin.after_agent_callback( + agent=self.mock_agent, callback_context=self.callback_context + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "AGENT_COMPLETED") + assert log_entry["content"] == "Agent Name: MyTestAgent" + + @pytest.mark.asyncio + async def test_before_model_callback_logs_correctly(self): + llm_request = llm_request_lib.LlmRequest( + model="gemini-pro", + contents=[types.Content(parts=[types.Part(text="Prompt")])], + config=types.GenerateContentConfig( + temperature=0.5, + top_p=0.9, + max_output_tokens=100, + system_instruction="Be helpful", + ), + tools_dict={ + "my_tool": mock.create_autospec( + base_tool_lib.BaseTool, instance=True + ) + }, # Fixed mock + ) + + await self.plugin.before_model_callback( + callback_context=self.callback_context, llm_request=llm_request + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "LLM_REQUEST") + assert "Model: gemini-pro" in log_entry["content"] + assert "System Prompt: Be helpful" in log_entry["content"] + assert ( + "Params: {temperature=0.5, top_p=0.9, max_output_tokens=100}" + in log_entry["content"] + ) + assert "Available Tools: ['my_tool']" in log_entry["content"] + + @pytest.mark.asyncio + async def test_after_model_callback_text_response(self): + llm_response = llm_response_lib.LlmResponse( + content=types.Content(parts=[types.Part(text="Model response")]), + usage_metadata=types.UsageMetadata( + prompt_token_count=10, + total_token_count=15, + ), + ) + await self.plugin.after_model_callback( + callback_context=self.callback_context, llm_response=llm_response + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "LLM_RESPONSE") + assert ( + "Tool Name: text_response, text: 'Model response'" + in log_entry["content"] + ) + # Adjusted assertion to expect None for candidates + assert "Token Usage: {prompt: 10" in log_entry["content"] + assert log_entry["error_message"] is None + + @pytest.mark.asyncio + async def test_after_model_callback_tool_call(self): + llm_response = llm_response_lib.LlmResponse( + content=types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall(name="tool1", args={}) + ) + ] + ), + ) + await self.plugin.after_model_callback( + callback_context=self.callback_context, llm_response=llm_response + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "LLM_RESPONSE") + assert "Tool Name: tool1" in log_entry["content"] + + @pytest.mark.asyncio + async def test_before_tool_callback_logs_correctly(self): + mock_tool = mock.create_autospec(base_tool_lib.BaseTool, instance=True) + mock_tool.name = "MyTool" + mock_tool.description = "Does something" + tool_args = {"param": "value"} + await self.plugin.before_tool_callback( + tool=mock_tool, tool_args=tool_args, tool_context=self.tool_context + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "TOOL_STARTING") + assert "Tool Name: MyTool" in log_entry["content"] + assert "Description: Does something" in log_entry["content"] + assert "Arguments: {'param': 'value'}" in log_entry["content"] + + @pytest.mark.asyncio + async def test_after_tool_callback_logs_correctly(self): + mock_tool = mock.create_autospec(base_tool_lib.BaseTool, instance=True) + mock_tool.name = "MyTool" + tool_args = {"param": "value"} + result = {"status": "success"} + await self.plugin.after_tool_callback( + tool=mock_tool, + tool_args=tool_args, + tool_context=self.tool_context, + result=result, + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "TOOL_COMPLETED") + assert "Tool Name: MyTool" in log_entry["content"] + assert "Result: {'status': 'success'}" in log_entry["content"] + + @pytest.mark.asyncio + async def test_on_model_error_callback_logs_correctly(self): + llm_request = mock.create_autospec( + llm_request_lib.LlmRequest, instance=True + ) + error = ValueError("LLM failed") + await self.plugin.on_model_error_callback( + callback_context=self.callback_context, + llm_request=llm_request, + error=error, + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "LLM_ERROR") + assert log_entry["content"] is None + assert log_entry["error_message"] == "LLM failed" + + @pytest.mark.asyncio + async def test_on_tool_error_callback_logs_correctly(self): + mock_tool = mock.create_autospec(base_tool_lib.BaseTool, instance=True) + mock_tool.name = "MyTool" + error = TimeoutError("Tool timed out") + await self.plugin.on_tool_error_callback( + tool=mock_tool, + tool_args={"param": "value"}, + tool_context=self.tool_context, + error=error, + ) + log_entry = self._get_logged_entry() + self._assert_common_fields(log_entry, "TOOL_ERROR") + assert log_entry["content"] == "Tool Name: MyTool" + assert log_entry["error_message"] == "Tool timed out"