diff --git a/src/google/adk/plugins/bigquery_logging_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py similarity index 56% rename from src/google/adk/plugins/bigquery_logging_plugin.py rename to src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 1191a6fe..cc1bf6a7 100644 --- a/src/google/adk/plugins/bigquery_logging_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -19,20 +19,16 @@ from datetime import datetime from datetime import timezone import json import logging -import threading from typing import Any from typing import Callable -from typing import Dict from typing import List from typing import Optional +from typing import Set from typing import TYPE_CHECKING -import warnings -import google.api_core.client_info +from google.api_core.gapic_v1 import client_info as gapic_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.cloud.bigquery import schema as bq_schema from google.cloud.bigquery_storage_v1 import types as bq_storage_types from google.cloud.bigquery_storage_v1.services.big_query_write.async_client import BigQueryWriteAsyncClient @@ -53,23 +49,29 @@ if TYPE_CHECKING: from ..agents.invocation_context import InvocationContext +# --- PyArrow Helper Functions --- def _pyarrow_datetime(): + """Returns PyArrow type for BigQuery DATETIME.""" return pa.timestamp("us", tz=None) def _pyarrow_numeric(): + """Returns PyArrow type for BigQuery NUMERIC.""" return pa.decimal128(38, 9) def _pyarrow_bignumeric(): + """Returns PyArrow type for BigQuery BIGNUMERIC.""" return pa.decimal256(76, 38) def _pyarrow_time(): + """Returns PyArrow type for BigQuery TIME.""" return pa.time64("us") def _pyarrow_timestamp(): + """Returns PyArrow type for BigQuery TIMESTAMP.""" return pa.timestamp("us", tz="UTC") @@ -92,11 +94,6 @@ _BQ_TO_ARROW_SCALARS = { "TIMESTAMP": _pyarrow_timestamp, } - -def _bq_to_arrow_scalars(bq_scalar: str): - return _BQ_TO_ARROW_SCALARS.get(bq_scalar) - - _BQ_FIELD_TYPE_TO_ARROW_FIELD_METADATA = { "GEOGRAPHY": { b"ARROW:extension:name": b"google:sqlType:geography", @@ -108,89 +105,114 @@ _BQ_FIELD_TYPE_TO_ARROW_FIELD_METADATA = { _STRUCT_TYPES = ("RECORD", "STRUCT") +def _bq_to_arrow_scalars(bq_scalar: str): + """Converts a BigQuery scalar type string to a PyArrow data type constructor.""" + return _BQ_TO_ARROW_SCALARS.get(bq_scalar) + + def _bq_to_arrow_struct_data_type(field): + """Converts a BigQuery STRUCT/RECORD field to a PyArrow struct type.""" arrow_fields = [] for subfield in field.fields: arrow_subfield = _bq_to_arrow_field(subfield) if arrow_subfield: arrow_fields.append(arrow_subfield) else: + logging.warning( + "Failed to convert STRUCT/RECORD field '%s' due to subfield '%s'.", + field.name, + subfield.name, + ) return None return pa.struct(arrow_fields) def _bq_to_arrow_range_data_type(field): + """Converts a BigQuery RANGE field to a PyArrow struct type.""" if field is None: raise ValueError("Range element type cannot be None") - element_type = field.element_type.upper() - arrow_element_type = _bq_to_arrow_scalars(element_type)() - return pa.struct([("start", arrow_element_type), ("end", arrow_element_type)]) + return pa.struct([ + ("start", _bq_to_arrow_scalars(field.element_type.upper())()), + ("end", _bq_to_arrow_scalars(field.element_type.upper())()), + ]) def _bq_to_arrow_data_type(field): - if field.mode is not None and field.mode.upper() == "REPEATED": - inner_type = _bq_to_arrow_data_type( - bq_schema.SchemaField(field.name, field.field_type, fields=field.fields) + """Converts a BigQuery schema field to a PyArrow data type.""" + if field.mode == "REPEATED": + inner = _bq_to_arrow_data_type( + bq_schema.SchemaField( + field.name, + field.field_type, + fields=field.fields, + range_element_type=getattr(field, "range_element_type", None), + ) ) - if inner_type: - return pa.list_(inner_type) - return None - + return pa.list_(inner) if inner else None field_type_upper = field.field_type.upper() if field.field_type else "" if field_type_upper in _STRUCT_TYPES: return _bq_to_arrow_struct_data_type(field) - if field_type_upper == "RANGE": return _bq_to_arrow_range_data_type(field.range_element_type) - - data_type_constructor = _bq_to_arrow_scalars(field_type_upper) - if data_type_constructor is None: + constructor = _bq_to_arrow_scalars(field_type_upper) + if constructor: + return constructor() + else: + logging.warning( + "Failed to convert BigQuery field '%s': unsupported type '%s'.", + field.name, + field.field_type, + ) return None - return data_type_constructor() -def _bq_to_arrow_field(bq_field, array_type=None): +def _bq_to_arrow_field(bq_field): + """Converts a BigQuery SchemaField to a PyArrow Field.""" arrow_type = _bq_to_arrow_data_type(bq_field) - if arrow_type is not None: - if array_type is not None: - arrow_type = array_type + if arrow_type: metadata = _BQ_FIELD_TYPE_TO_ARROW_FIELD_METADATA.get( bq_field.field_type.upper() if bq_field.field_type else "" ) return pa.field( bq_field.name, arrow_type, - nullable=False if bq_field.mode.upper() == "REPEATED" else True, + nullable=(bq_field.mode != "REPEATED"), metadata=metadata, ) - - warnings.warn(f"Unable to determine Arrow type for field '{bq_field.name}'.") + logging.warning( + "Could not determine Arrow type for field '%s' with type '%s'.", + bq_field.name, + bq_field.field_type, + ) return None def to_arrow_schema(bq_schema_list): - """Return the Arrow schema, corresponding to a given BigQuery schema.""" + """Converts a list of BigQuery SchemaFields to a PyArrow Schema.""" arrow_fields = [] for bq_field in bq_schema_list: - arrow_field = _bq_to_arrow_field(bq_field) - if arrow_field is None: + af = _bq_to_arrow_field(bq_field) + if af: + arrow_fields.append(af) + else: + logging.warning( + "Failed to convert schema due to field '%s'.", bq_field.name + ) return None - arrow_fields.append(arrow_field) return pa.schema(arrow_fields) @dataclasses.dataclass class BigQueryLoggerConfig: - """Configuration for the BigQueryAgentAnalyticsPlugin. + """Configuration for BigQueryAgentAnalyticsPlugin. Attributes: - enabled: Whether the plugin is enabled. - event_allowlist: List of event types to log. If None, all are allowed - except those in event_denylist. - event_denylist: List of event types to not log. Takes precedence over - event_allowlist. - content_formatter: Function to format or redact the 'content' field before - logging. + enabled: Whether logging is enabled. + event_allowlist: A list of event types to log. If None, all events are + logged except those in event_denylist. + event_denylist: A list of event types to skip logging. + content_formatter: An optional function to format event content before + logging. """ enabled: bool = True @@ -199,7 +221,9 @@ class BigQueryLoggerConfig: content_formatter: Optional[Callable[[Any], str]] = None +# --- Helper Formatters --- def _get_event_type(event: Event) -> str: + """Determines the event type from an Event object.""" if event.author == "user": return "USER_INPUT" if event.get_function_calls(): @@ -210,59 +234,55 @@ def _get_event_type(event: Event) -> str: return "MODEL_RESPONSE" if event.error_message: return "ERROR" - return "SYSTEM" # Fallback for other event types + return "SYSTEM" def _format_content( - content: Optional[types.Content], max_length: int = 200 + content: Optional[types.Content], max_len: int = 500 ) -> str: - """Format content for logging, truncating if too long.""" + """Formats an Event content for logging.""" 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") + for p in content.parts: + if p.text: + parts.append( + f"text: '{p.text[:max_len]}...' " + if len(p.text) > max_len + else f"text: '{p.text}'" + ) + elif p.function_call: + parts.append(f"call: {p.function_call.name}") + elif p.function_response: + parts.append(f"resp: {p.function_response.name}") else: - parts.append("other_part") + parts.append("other") return " | ".join(parts) -def _format_args(args: dict[str, Any], max_length: int = 300) -> str: - """Format arguments dictionary for logging.""" +def _format_args(args: dict[str, Any], max_len: int = 1000) -> str: + """Formats tool arguments or results for logging.""" if not args: return "{}" - formatted = str(args) - if len(formatted) > max_length: - formatted = formatted[:max_length] + "...}" - return formatted + try: + s = json.dumps(args) + except TypeError: + s = str(args) + return s[:max_len] + "..." if len(s) > max_len else s class BigQueryAgentAnalyticsPlugin(BasePlugin): - """A plugin that logs ADK events to a BigQuery table. + """A plugin that logs agent analytic events to Google BigQuery. - 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. + This plugin captures key events during an agent's lifecycle—such as user + interactions, tool executions, LLM requests/responses, and errors—and + streams them to a BigQuery table for analysis and monitoring. - 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 - - Logging behavior can be customized using the BigQueryLoggerConfig. + It uses the BigQuery Write API for efficient, high-throughput streaming + ingestion and is designed to be non-blocking, ensuring that logging + operations do not impact agent performance. If the destination table does + not exist, the plugin will attempt to create it based on a predefined + schema. """ def __init__( @@ -273,237 +293,231 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): config: Optional[BigQueryLoggerConfig] = None, **kwargs, ): + """Initializes the BigQueryAgentAnalyticsPlugin. + + Args: + project_id: Google Cloud project ID. + dataset_id: BigQuery dataset ID. + table_id: BigQuery table ID for agent events. + config: Plugin configuration. + **kwargs: Additional arguments. + """ super().__init__(name=kwargs.get("name", "BigQueryAgentAnalyticsPlugin")) - self._project_id = project_id - self._dataset_id = dataset_id - self._table_id = table_id + self._project_id, self._dataset_id, self._table_id = ( + project_id, + dataset_id, + table_id, + ) self._config = config if config else BigQueryLoggerConfig() self._bq_client: bigquery.Client | None = None - self._client_init_lock = threading.Lock() - self._init_done = False - self._init_succeeded = False - self._write_client: BigQueryWriteAsyncClient | None = None + self._init_lock: asyncio.Lock | None = None self._arrow_schema: pa.Schema | None = None - if not self._config.enabled: - logging.info( - "BigQueryAgentAnalyticsPlugin %s is disabled by configuration.", - self.name, - ) - return + self._background_tasks: Set[asyncio.Task] = set() # Track pending logs + self._schema = [ + 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"), + ] - logging.debug( - "DEBUG: BigQueryAgentAnalyticsPlugin INSTANTIATED (Name: %s)", self.name - ) + def _format_content_safely( + self, content: Optional[types.Content] + ) -> str | None: + """Formats content using self._config.content_formatter or _format_content, catching errors.""" + if content is None: + return None + try: + if self._config.content_formatter: + return self._config.content_formatter(content) + return _format_content(content) + except Exception as e: + logging.warning(f"Content formatter failed: {e}") + return "[FORMATTING FAILED]" - def _ensure_initialized_sync(self): - """Synchronous initialization of BQ client and table.""" - if not self._config.enabled: - return - - with self._client_init_lock: - if self._init_done: - return - self._init_done = True + async def _ensure_init(self): + """Ensures BigQuery clients are initialized.""" + if self._write_client: + return True + if not self._init_lock: + self._init_lock = asyncio.Lock() + async with self._init_lock: + if self._write_client: + return True try: - credentials, _ = google.auth.default( - scopes=[ - "https://www.googleapis.com/auth/bigquery", - "https://www.googleapis.com/auth/cloud-platform", # For Storage Write - ] + creds, _ = await asyncio.to_thread( + google.auth.default, + scopes=["https://www.googleapis.com/auth/cloud-platform"], ) - client_info = google.api_core.client_info.ClientInfo( + client_info = gapic_client_info.ClientInfo( user_agent=f"google-adk-bq-logger/{version.__version__}" ) - - # 1. Init BQ Client (for create_dataset/create_table) self._bq_client = bigquery.Client( - project=self._project_id, - credentials=credentials, + project=self._project_id, credentials=creds, client_info=client_info + ) + + # Ensure table exists (sync call in thread) + def create_resources(): + if self._bq_client: + dataset = self._bq_client.create_dataset( + self._dataset_id, exists_ok=True + ) + table = bigquery.Table( + f"{self._project_id}.{self._dataset_id}.{self._table_id}", + schema=self._schema, + ) + self._bq_client.create_table(table, exists_ok=True) + + await asyncio.to_thread(create_resources) + + self._write_client = BigQueryWriteAsyncClient( + credentials=creds, client_info=client_info, ) + self._arrow_schema = to_arrow_schema(self._schema) + return True + except Exception as e: + logging.error(f"BQ Init Failed: {e}") + return False - # 2. Init BQ Storage Write Client - self._write_client = BigQueryWriteAsyncClient( - credentials=credentials, client_info=client_info - ) + async def _perform_write(self, row: dict): + """Actual async write operation, intended to run as a background task.""" + try: + if ( + not await self._ensure_init() + or not self._write_client + or not self._arrow_schema + ): + return - logging.info( - "BigQuery clients (Core & Storage Write) 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) + # Serialize + pydict = {f.name: [row.get(f.name)] for f in self._arrow_schema} + batch = pa.RecordBatch.from_pydict(pydict, schema=self._arrow_schema) + req = bq_storage_types.AppendRowsRequest( + write_stream=f"projects/{self._project_id}/datasets/{self._dataset_id}/tables/{self._table_id}/_default" + ) + req.arrow_rows.writer_schema.serialized_schema = ( + self._arrow_schema.serialize().to_pybytes() + ) + req.arrow_rows.rows.serialized_record_batch = ( + batch.serialize().to_pybytes() + ) - # Schema - schema = [ - 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) + # Write with protection against immediate cancellation + async for resp in await asyncio.shield( + self._write_client.append_rows(iter([req])) + ): + if resp.error.code != 0: + logging.error(f"BQ Write Error: {resp.error.message}") - # 4. Store Arrow schema for Write API - self._arrow_schema = to_arrow_schema(schema) # USE LOCAL VERSION - # --- self._table_ref_str removed --- + except RuntimeError as e: + # Silently ignore event loop closed errors during background writes + if "Event loop is closed" not in str(e): + logging.exception(f"BQ Runtime Error: {e}") + except Exception as e: + logging.error(f"BQ Write Failed: {e}") - 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]): + async def _log(self, data: dict): + """Schedules a log entry to be written in the background.""" if not self._config.enabled: return - - event_type = event_dict.get("event_type") - - # Check denylist + event_type = data.get("event_type") if ( self._config.event_denylist and event_type in self._config.event_denylist ): return - - # Check allowlist if ( self._config.event_allowlist and event_type not in self._config.event_allowlist ): return - # Apply custom content formatter - if self._config.content_formatter and "content" in event_dict: + # Prepare row immediately (capture current state) + row = { + "timestamp": datetime.now(timezone.utc), + "event_type": None, + "agent": None, + "session_id": None, + "invocation_id": None, + "user_id": None, + "content": None, + "error_message": None, + } + row.update(data) + + # Fire and forget: Create task and track it + task = asyncio.create_task(self._perform_write(row)) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + + async def shutdown(self): + """Flushes pending logs and closes client.""" + # 1. Wait for pending background logs (best effort, 2s timeout) + if self._background_tasks: + logging.info(f"Flushing {len(self._background_tasks)} pending BQ logs...") + done, pending = await asyncio.wait(self._background_tasks, timeout=2.0) + if pending: + logging.warning( + f"{len(pending)} BQ logs could not be flushed before shutdown." + ) + + # 2. Close client + if self._write_client and self._write_client.transport: try: - event_dict["content"] = self._config.content_formatter( - event_dict["content"] + logging.info("Closing BQ Write client transport...") + await asyncio.wait_for( + self._write_client.transport.close(), timeout=1.0 ) except Exception as e: - logging.warning( - "Error applying custom content formatter for event type %s: %s", - event_type, - e, - ) - # Optionally log a generic message or the error - - try: - if not self._init_done: - await asyncio.to_thread(self._ensure_initialized_sync) - - # Check for all required Storage Write API components - if not ( - self._init_succeeded and self._write_client and self._arrow_schema - ): - logging.warning("BigQuery write client not initialized. Skipping log.") - return - - default_row = { - "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} - - # --- START MODIFIED STORAGE WRITE API LOGIC (using Default Stream) --- - # 1. Convert the single row dict to a PyArrow RecordBatch - # pa.RecordBatch.from_pydict requires a dict of lists - pydict = { - field.name: [insert_row.get(field.name)] - for field in self._arrow_schema - } - batch = pa.RecordBatch.from_pydict(pydict, schema=self._arrow_schema) - - # 2. Create the AppendRowsRequest, pointing to the default stream - request = bq_storage_types.AppendRowsRequest( - write_stream=( - f"projects/{self._project_id}/datasets/{self._dataset_id}" - f"/tables/{self._table_id}/_default" - ) - ) - request.arrow_rows.writer_schema.serialized_schema = ( - self._arrow_schema.serialize().to_pybytes() - ) - - request.arrow_rows.rows.serialized_record_batch = ( - batch.serialize().to_pybytes() - ) - - # 3. Send the request and check for errors - response_iterator = self._write_client.append_rows(requests=[request]) - async for response in response_iterator: - if response.row_errors: - logging.error( - "Errors occurred while writing to BigQuery (Storage Write" - " API): %s", - response.row_errors, - ) - break # Only one response expected - except Exception as e: - logging.exception("Failed to log to BigQuery: %s", e) + logging.warning(f"Error during BQ Write client transport close: {e}") + self._write_client = None + if self._bq_client: + try: + logging.info("Closing BQ client...") + self._bq_client.close() + except Exception as e: + logging.warning(f"Error during BQ client close: {e}") + self._bq_client = None + # --- Streamlined Callbacks --- 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 = { - "timestamp": datetime.now(timezone.utc).isoformat(), + ) -> None: + """Callback for user messages.""" + 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: {_format_content(user_message)}", - } - await self._log_to_bigquery_async(event_dict) - return None + "content": f"User Content: {self._format_content_safely(user_message)}", + }) async def before_run_callback( self, *, invocation_context: InvocationContext - ) -> Optional[types.Content]: - """Log invocation start.""" - event_dict = { - "timestamp": datetime.now(timezone.utc).isoformat(), + ) -> None: + """Callback before agent invocation.""" + await self._log({ "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, - "content": None, - } - 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 = { - "timestamp": datetime.fromtimestamp( - event.timestamp, timezone.utc - ).isoformat(), + ) -> None: + """Callback for agent events.""" + await self._log({ "event_type": _get_event_type(event), "agent": event.author, "session_id": invocation_context.session.id, @@ -517,68 +531,54 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): else None ), "error_message": event.error_message, - } - await self._log_to_bigquery_async(event_dict) - return None + "timestamp": datetime.fromtimestamp(event.timestamp, timezone.utc), + }) async def after_run_callback( self, *, invocation_context: InvocationContext - ) -> Optional[None]: - """Log invocation completion.""" - event_dict = { - "timestamp": datetime.now(timezone.utc).isoformat(), + ) -> None: + """Callback after agent invocation.""" + await self._log({ "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, - "content": None, - } - 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 = { - "timestamp": datetime.now(timezone.utc).isoformat(), + ) -> None: + """Callback before an agent starts.""" + await self._log({ "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 = { - "timestamp": datetime.now(timezone.utc).isoformat(), + ) -> None: + """Callback after an agent completes.""" + await self._log({ "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.""" - + ) -> None: + """Callback before LLM call.""" content_parts = [ f"Model: {llm_request.model or 'default'}", ] - - # Log Full System Instruction system_instruction_text = "None" if llm_request.config and llm_request.config.system_instruction: si = llm_request.config.system_instruction @@ -602,8 +602,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): 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 = {} @@ -629,23 +627,19 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): ) final_content = " | ".join(content_parts) - - event_dict = { - "timestamp": datetime.now(timezone.utc).isoformat(), + await self._log({ "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.""" + ) -> None: + """Callback after LLM call.""" content_parts = [] content = llm_response.content is_tool_call = False @@ -653,7 +647,6 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): 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 = [ @@ -663,10 +656,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): ] 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...'" + text_content = self._format_content_safely(llm_response.content) content_parts.append(f"Tool Name: text_response, {text_content}") if llm_response.usage_metadata: @@ -686,21 +676,15 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): content_parts.append(token_usage_str) final_content = " | ".join(content_parts) - - event_dict = { - "timestamp": datetime.now(timezone.utc).isoformat(), + await self._log({ "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 + "error_message": llm_response.error_message, + }) async def before_tool_callback( self, @@ -708,10 +692,9 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): tool: BaseTool, tool_args: dict[str, Any], tool_context: ToolContext, - ) -> Optional[None]: - """Log tool execution start.""" - event_dict = { - "timestamp": datetime.now(timezone.utc).isoformat(), + ) -> None: + """Callback before tool call.""" + await self._log({ "event_type": "TOOL_STARTING", "agent": tool_context.agent_name, "session_id": tool_context.session.id, @@ -721,9 +704,7 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): 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, @@ -733,18 +714,15 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): tool_context: ToolContext, result: dict[str, Any], ) -> None: - """Log tool execution completion.""" - event_dict = { - "timestamp": datetime.now(timezone.utc).isoformat(), + """Callback after tool call.""" + 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: {_format_args(result)}", - } - await self._log_to_bigquery_async(event_dict) - return None + }) async def on_model_error_callback( self, @@ -752,19 +730,16 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): callback_context: CallbackContext, llm_request: LlmRequest, error: Exception, - ) -> Optional[LlmResponse]: - """Log LLM error.""" - event_dict = { - "timestamp": datetime.now(timezone.utc).isoformat(), + ) -> None: + """Callback for LLM errors.""" + await self._log({ "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, @@ -774,9 +749,8 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): tool_context: ToolContext, error: Exception, ) -> None: - """Log tool error.""" - event_dict = { - "timestamp": datetime.now(timezone.utc).isoformat(), + """Callback for tool errors.""" + await self._log({ "event_type": "TOOL_ERROR", "agent": tool_context.agent_name, "session_id": tool_context.session.id, @@ -786,6 +760,4 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): f"Tool Name: {tool.name}, Arguments: {_format_args(tool_args)}" ), "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_agent_analytics_plugin.py similarity index 75% rename from tests/unittests/plugins/test_bigquery_logging_plugin.py rename to tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index ae9a1f26..f3251119 100644 --- a/tests/unittests/plugins/test_bigquery_logging_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -12,8 +12,7 @@ # 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 @@ -25,7 +24,7 @@ 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 bigquery_agent_analytics_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 @@ -35,11 +34,12 @@ import google.auth from google.auth import exceptions as auth_exceptions import google.auth.credentials from google.cloud import bigquery +from google.cloud.bigquery_storage_v1 import types as bq_storage_types from google.genai import types import pyarrow as pa import pytest -BigQueryLoggerConfig = bigquery_logging_plugin.BigQueryLoggerConfig +BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig PROJECT_ID = "test-gcp-project" DATASET_ID = "adk_logs" @@ -125,25 +125,28 @@ def mock_bq_client(): @pytest.fixture def mock_write_client(): with mock.patch.object( - bigquery_logging_plugin, "BigQueryWriteAsyncClient", autospec=True + bigquery_agent_analytics_plugin, "BigQueryWriteAsyncClient", autospec=True ) as mock_cls: mock_client = mock_cls.return_value - mock_append_rows_response = mock.MagicMock() - # Configure the 'row_errors' attribute on the mock object. - mock_append_rows_response.row_errors = [] - mock_append_rows_response.error = mock.MagicMock() - mock_append_rows_response.error.code = 0 # OK status + mock_client.transport = mock.AsyncMock() - mock_client.append_rows.return_value = _async_gen(mock_append_rows_response) + async def fake_append_rows(requests, **kwargs): + # This function is now async, so `await client.append_rows` works. + mock_append_rows_response = mock.MagicMock() + mock_append_rows_response.row_errors = [] + mock_append_rows_response.error = mock.MagicMock() + mock_append_rows_response.error.code = 0 # OK status + # This a gen is what's returned *after* the await. + return _async_gen(mock_append_rows_response) + + mock_client.append_rows.side_effect = fake_append_rows yield mock_client @pytest.fixture def dummy_arrow_schema(): return pa.schema([ - pa.field( - "timestamp", pa.string() - ), # Store as string for simplicity in test + pa.field("timestamp", pa.timestamp("us", tz="UTC")), pa.field("event_type", pa.string()), pa.field("agent", pa.string()), pa.field("session_id", pa.string()), @@ -157,7 +160,7 @@ def dummy_arrow_schema(): @pytest.fixture def mock_to_arrow_schema(dummy_arrow_schema): with mock.patch.object( - bigquery_logging_plugin, + bigquery_agent_analytics_plugin, "to_arrow_schema", autospec=True, return_value=dummy_arrow_schema, @@ -177,19 +180,19 @@ def mock_asyncio_to_thread(): @pytest.fixture -def bq_plugin_inst( +async def bq_plugin_inst( mock_auth_default, mock_bq_client, mock_write_client, mock_to_arrow_schema, + mock_asyncio_to_thread, ): - plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin( + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( project_id=PROJECT_ID, dataset_id=DATASET_ID, table_id=TABLE_ID, ) - # Trigger lazy initialization - plugin._ensure_initialized_sync() + await plugin._ensure_init() # Ensure clients are initialized mock_write_client.append_rows.reset_mock() return plugin @@ -205,7 +208,8 @@ def _get_captured_event_dict(mock_write_client, expected_schema): """Helper to get the event_dict passed to append_rows.""" mock_write_client.append_rows.assert_called_once() call_args = mock_write_client.append_rows.call_args - requests = call_args.kwargs["requests"] + requests_iter = call_args.args[0] + requests = list(requests_iter) assert len(requests) == 1 request = requests[0] assert request.write_stream == DEFAULT_STREAM_NAME @@ -228,7 +232,7 @@ def _assert_common_fields(log_entry, event_type, agent="MyTestAgent"): assert log_entry["invocation_id"] == "inv-789" assert log_entry["user_id"] == "user-456" assert "timestamp" in log_entry - assert isinstance(log_entry["timestamp"], str) + assert isinstance(log_entry["timestamp"], datetime.datetime) # --- Test Class --- @@ -245,17 +249,17 @@ class TestBigQueryAgentAnalyticsPlugin: invocation_context, ): config = BigQueryLoggerConfig(enabled=False) - plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin( + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( project_id=PROJECT_ID, dataset_id=DATASET_ID, table_id=TABLE_ID, config=config, ) - plugin._ensure_initialized_sync() # Should do nothing - user_message = types.Content(parts=[types.Part(text="Test")]) + # user_message = types.Content(parts=[types.Part(text="Test")]) await plugin.on_user_message_callback( - invocation_context=invocation_context, user_message=user_message + invocation_context=invocation_context, + user_message=types.Content(parts=[types.Part(text="Test")]), ) mock_auth_default.assert_not_called() mock_bq_client.assert_not_called() @@ -271,12 +275,13 @@ class TestBigQueryAgentAnalyticsPlugin: mock_bq_client, mock_to_arrow_schema, dummy_arrow_schema, + mock_asyncio_to_thread, ): config = BigQueryLoggerConfig(event_allowlist=["LLM_REQUEST"]) - plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin( + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( PROJECT_ID, DATASET_ID, TABLE_ID, config ) - plugin._ensure_initialized_sync() + await plugin._ensure_init() mock_write_client.append_rows.reset_mock() llm_request = llm_request_lib.LlmRequest( @@ -286,6 +291,7 @@ class TestBigQueryAgentAnalyticsPlugin: await plugin.before_model_callback( callback_context=callback_context, llm_request=llm_request ) + await asyncio.sleep(0.01) # Allow background task to run mock_write_client.append_rows.assert_called_once() mock_write_client.append_rows.reset_mock() @@ -293,6 +299,7 @@ class TestBigQueryAgentAnalyticsPlugin: await plugin.on_user_message_callback( invocation_context=invocation_context, user_message=user_message ) + await asyncio.sleep(0.01) # Allow background task to run mock_write_client.append_rows.assert_not_called() @pytest.mark.asyncio @@ -304,21 +311,24 @@ class TestBigQueryAgentAnalyticsPlugin: mock_bq_client, mock_to_arrow_schema, dummy_arrow_schema, + mock_asyncio_to_thread, ): config = BigQueryLoggerConfig(event_denylist=["USER_MESSAGE_RECEIVED"]) - plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin( + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( PROJECT_ID, DATASET_ID, TABLE_ID, config ) - plugin._ensure_initialized_sync() + await plugin._ensure_init() mock_write_client.append_rows.reset_mock() user_message = types.Content(parts=[types.Part(text="What is up?")]) await plugin.on_user_message_callback( invocation_context=invocation_context, user_message=user_message ) + await asyncio.sleep(0.01) mock_write_client.append_rows.assert_not_called() await plugin.before_run_callback(invocation_context=invocation_context) + await asyncio.sleep(0.01) mock_write_client.append_rows.assert_called_once() @pytest.mark.asyncio @@ -330,24 +340,26 @@ class TestBigQueryAgentAnalyticsPlugin: mock_bq_client, mock_to_arrow_schema, dummy_arrow_schema, + mock_asyncio_to_thread, ): def redact_content(content): return "[REDACTED]" config = BigQueryLoggerConfig(content_formatter=redact_content) - plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin( + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( PROJECT_ID, DATASET_ID, TABLE_ID, config ) - plugin._ensure_initialized_sync() + await plugin._ensure_init() mock_write_client.append_rows.reset_mock() user_message = types.Content(parts=[types.Part(text="Secret message")]) await plugin.on_user_message_callback( invocation_context=invocation_context, user_message=user_message ) + 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_common_fields(log_entry, "USER_MESSAGE_RECEIVED") - assert log_entry["content"] == "[REDACTED]" + assert log_entry["content"] == "User Content: [REDACTED]" @pytest.mark.asyncio async def test_content_formatter_error( @@ -358,31 +370,26 @@ class TestBigQueryAgentAnalyticsPlugin: mock_bq_client, mock_to_arrow_schema, dummy_arrow_schema, + mock_asyncio_to_thread, ): def error_formatter(content): raise ValueError("Formatter failed") config = BigQueryLoggerConfig(content_formatter=error_formatter) - plugin = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin( + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( PROJECT_ID, DATASET_ID, TABLE_ID, config ) - plugin._ensure_initialized_sync() + await plugin._ensure_init() mock_write_client.append_rows.reset_mock() - user_message = types.Content(parts=[types.Part(text="Test")]) - - with mock.patch.object(logging, "warning") as mock_log_warning: - await plugin.on_user_message_callback( - invocation_context=invocation_context, user_message=user_message - ) - mock_log_warning.assert_called_once_with( - "Error applying custom content formatter for event type %s: %s", - "USER_MESSAGE_RECEIVED", - mock.ANY, - ) + user_message = types.Content(parts=[types.Part(text="Secret message")]) + await plugin.on_user_message_callback( + invocation_context=invocation_context, user_message=user_message + ) + 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_common_fields(log_entry, "USER_MESSAGE_RECEIVED") - assert "User Content: text: 'Test'" in log_entry["content"] + assert log_entry["content"] == "User Content: [FORMATTING FAILED]" @pytest.mark.asyncio async def test_on_user_message_callback_logs_correctly( @@ -396,6 +403,7 @@ class TestBigQueryAgentAnalyticsPlugin: await bq_plugin_inst.on_user_message_callback( invocation_context=invocation_context, user_message=user_message ) + await asyncio.sleep(0.01) 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?'" @@ -419,12 +427,13 @@ class TestBigQueryAgentAnalyticsPlugin: await bq_plugin_inst.on_event_callback( invocation_context=invocation_context, event=event ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) - _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" + _assert_common_fields(log_entry, "TOOL_CALL", agent="MyTestAgent") + assert '"name": "get_weather"' in log_entry["content"] + assert log_entry["timestamp"] == datetime.datetime( + 2025, 10, 22, 10, 0, 0, tzinfo=datetime.timezone.utc + ) @pytest.mark.asyncio async def test_on_event_callback_model_response( @@ -444,57 +453,74 @@ class TestBigQueryAgentAnalyticsPlugin: await bq_plugin_inst.on_event_callback( invocation_context=invocation_context, event=event ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) - _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" + _assert_common_fields(log_entry, "MODEL_RESPONSE", agent="MyTestAgent") + assert '"text": "Hello there!"' in log_entry["content"] + assert log_entry["timestamp"] == datetime.datetime( + 2025, 10, 22, 11, 0, 0, tzinfo=datetime.timezone.utc + ) @pytest.mark.asyncio async def test_bigquery_client_initialization_failure( - self, mock_auth_default, mock_write_client, invocation_context + self, + mock_auth_default, + mock_write_client, + invocation_context, + mock_asyncio_to_thread, ): mock_auth_default.side_effect = auth_exceptions.GoogleAuthError( "Auth failed" ) - plugin_with_fail = bigquery_logging_plugin.BigQueryAgentAnalyticsPlugin( - project_id=PROJECT_ID, - dataset_id=DATASET_ID, - table_id=TABLE_ID, + plugin_with_fail = ( + bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + project_id=PROJECT_ID, + dataset_id=DATASET_ID, + table_id=TABLE_ID, + ) ) - with mock.patch.object(logging, "exception") as mock_log_exception: + with mock.patch.object(logging, "error") as mock_log_error: await plugin_with_fail.on_user_message_callback( invocation_context=invocation_context, user_message=types.Content(parts=[types.Part(text="Test")]), ) - mock_log_exception.assert_called_once_with( - "Failed to initialize BigQuery client or table: %s", mock.ANY - ) + await asyncio.sleep(0.01) + mock_log_error.assert_any_call("BQ Init Failed: Auth failed") mock_write_client.append_rows.assert_not_called() @pytest.mark.asyncio async def test_bigquery_insert_error_does_not_raise( self, bq_plugin_inst, mock_write_client, invocation_context ): - mock_append_rows_response = mock.MagicMock() - mock_append_rows_response.row_errors = [mock.MagicMock()] - mock_append_rows_response.error = mock.MagicMock() - mock_append_rows_response.error.code = 0 - mock_write_client.append_rows.return_value = _async_gen( - mock_append_rows_response - ) + + async def fake_append_rows_with_error(requests, **kwargs): + mock_append_rows_response = mock.MagicMock() + mock_append_rows_response.row_errors = [] # No row errors + mock_append_rows_response.error = mock.MagicMock() + mock_append_rows_response.error.code = 3 # INVALID_ARGUMENT + mock_append_rows_response.error.message = "Test BQ Error" + return _async_gen(mock_append_rows_response) + + mock_write_client.append_rows.side_effect = fake_append_rows_with_error with mock.patch.object(logging, "error") as mock_log_error: await bq_plugin_inst.on_user_message_callback( invocation_context=invocation_context, user_message=types.Content(parts=[types.Part(text="Test")]), ) - mock_log_error.assert_called_with( - "Errors occurred while writing to BigQuery (Storage Write API): %s", - mock_append_rows_response.row_errors, - ) + await asyncio.sleep(0.01) + mock_log_error.assert_called_with("BQ Write Error: Test BQ Error") mock_write_client.append_rows.assert_called_once() + @pytest.mark.asyncio + async def test_shutdown( + self, bq_plugin_inst, mock_bq_client, mock_write_client + ): + await bq_plugin_inst.shutdown() + mock_write_client.transport.close.assert_called_once() + mock_bq_client.close.assert_called_once() + + # ... other tests remain the same ... @pytest.mark.asyncio async def test_before_run_callback_logs_correctly( self, @@ -506,6 +532,7 @@ class TestBigQueryAgentAnalyticsPlugin: await bq_plugin_inst.before_run_callback( invocation_context=invocation_context ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "INVOCATION_STARTING") assert log_entry["content"] is None @@ -521,6 +548,7 @@ class TestBigQueryAgentAnalyticsPlugin: await bq_plugin_inst.after_run_callback( invocation_context=invocation_context ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "INVOCATION_COMPLETED") assert log_entry["content"] is None @@ -537,6 +565,7 @@ class TestBigQueryAgentAnalyticsPlugin: await bq_plugin_inst.before_agent_callback( agent=mock_agent, callback_context=callback_context ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "AGENT_STARTING") assert log_entry["content"] == "Agent Name: MyTestAgent" @@ -553,6 +582,7 @@ class TestBigQueryAgentAnalyticsPlugin: await bq_plugin_inst.after_agent_callback( agent=mock_agent, callback_context=callback_context ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "AGENT_COMPLETED") assert log_entry["content"] == "Agent Name: MyTestAgent" @@ -568,32 +598,14 @@ class TestBigQueryAgentAnalyticsPlugin: 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=types.Content( - parts=[types.Part(text="Be helpful")] - ), - ), - tools_dict={ - "my_tool": mock.create_autospec( - base_tool_lib.BaseTool, instance=True, spec_set=True - ) - }, ) await bq_plugin_inst.before_model_callback( callback_context=callback_context, llm_request=llm_request ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _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"] + assert log_entry["content"] == "Model: gemini-pro | System Prompt: Empty" @pytest.mark.asyncio async def test_after_model_callback_text_response( @@ -612,13 +624,16 @@ class TestBigQueryAgentAnalyticsPlugin: await bq_plugin_inst.after_model_callback( callback_context=callback_context, llm_response=llm_response ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "LLM_RESPONSE") assert ( "Tool Name: text_response, text: 'Model response'" in log_entry["content"] ) - assert "Token Usage: {prompt: 10," in log_entry["content"] + assert "Token Usage:" in log_entry["content"] + assert "prompt: 10" in log_entry["content"] + assert "total: 15" in log_entry["content"] assert log_entry["error_message"] is None @pytest.mark.asyncio @@ -629,21 +644,24 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, dummy_arrow_schema, ): + tool_fc = types.FunctionCall(name="get_weather", args={"location": "Paris"}) llm_response = llm_response_lib.LlmResponse( - content=types.Content( - parts=[ - types.Part( - function_call=types.FunctionCall(name="tool1", args={}) - ) - ] + content=types.Content(parts=[types.Part(function_call=tool_fc)]), + usage_metadata=types.UsageMetadata( + prompt_token_count=10, total_token_count=15 ), ) await bq_plugin_inst.after_model_callback( callback_context=callback_context, llm_response=llm_response ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "LLM_RESPONSE") - assert "Tool Name: tool1" in log_entry["content"] + assert "Tool Name: get_weather" in log_entry["content"] + assert "Token Usage:" in log_entry["content"] + assert "prompt: 10" in log_entry["content"] + assert "total: 15" in log_entry["content"] + assert log_entry["error_message"] is None @pytest.mark.asyncio async def test_before_tool_callback_logs_correctly( @@ -653,16 +671,18 @@ class TestBigQueryAgentAnalyticsPlugin: base_tool_lib.BaseTool, instance=True, spec_set=True ) type(mock_tool).name = mock.PropertyMock(return_value="MyTool") - type(mock_tool).description = mock.PropertyMock( - return_value="Does something" - ) + type(mock_tool).description = mock.PropertyMock(return_value="Description") await bq_plugin_inst.before_tool_callback( tool=mock_tool, tool_args={"param": "value"}, tool_context=tool_context ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "TOOL_STARTING") - assert "Tool Name: MyTool" in log_entry["content"] - assert "Arguments: {'param': 'value'}" in log_entry["content"] + assert ( + log_entry["content"] + == 'Tool Name: MyTool, Description: Description, Arguments: {"param":' + ' "value"}' + ) @pytest.mark.asyncio async def test_after_tool_callback_logs_correctly( @@ -672,16 +692,20 @@ class TestBigQueryAgentAnalyticsPlugin: base_tool_lib.BaseTool, instance=True, spec_set=True ) type(mock_tool).name = mock.PropertyMock(return_value="MyTool") + type(mock_tool).description = mock.PropertyMock(return_value="Description") await bq_plugin_inst.after_tool_callback( tool=mock_tool, tool_args={}, tool_context=tool_context, result={"status": "success"}, ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "TOOL_COMPLETED") - assert "Tool Name: MyTool" in log_entry["content"] - assert "Result: {'status': 'success'}" in log_entry["content"] + assert ( + log_entry["content"] + == 'Tool Name: MyTool, Result: {"status": "success"}' + ) @pytest.mark.asyncio async def test_on_model_error_callback_logs_correctly( @@ -699,6 +723,7 @@ class TestBigQueryAgentAnalyticsPlugin: await bq_plugin_inst.on_model_error_callback( callback_context=callback_context, llm_request=llm_request, error=error ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "LLM_ERROR") assert log_entry["content"] is None @@ -712,6 +737,7 @@ class TestBigQueryAgentAnalyticsPlugin: base_tool_lib.BaseTool, instance=True, spec_set=True ) type(mock_tool).name = mock.PropertyMock(return_value="MyTool") + type(mock_tool).description = mock.PropertyMock(return_value="Description") error = TimeoutError("Tool timed out") await bq_plugin_inst.on_tool_error_callback( tool=mock_tool, @@ -719,8 +745,11 @@ class TestBigQueryAgentAnalyticsPlugin: tool_context=tool_context, error=error, ) + await asyncio.sleep(0.01) log_entry = _get_captured_event_dict(mock_write_client, dummy_arrow_schema) _assert_common_fields(log_entry, "TOOL_ERROR") - assert "Tool Name: MyTool" in log_entry["content"] - assert "Arguments: {'param': 'value'}" in log_entry["content"] + assert ( + log_entry["content"] + == 'Tool Name: MyTool, Arguments: {"param": "value"}' + ) assert log_entry["error_message"] == "Tool timed out"