diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 36a23378..7cbf931c 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -31,6 +31,7 @@ import random import time from types import MappingProxyType from typing import Any +from typing import Awaitable from typing import Callable from typing import Optional from typing import TYPE_CHECKING @@ -409,6 +410,11 @@ class BigQueryLoggerConfig: # Format: "location.connection_id" (e.g. "us.my-connection") connection_id: Optional[str] = None + # Toggle for session metadata (e.g. gchat thread-id) + log_session_metadata: bool = True + # Static custom tags (e.g. {"agent_role": "sales"}) + custom_tags: dict[str, Any] = field(default_factory=dict) + # ============================================================================== # HELPER: TRACE MANAGER (Async-Safe with ContextVars) @@ -435,6 +441,9 @@ _span_id_stack_ctx: contextvars.ContextVar[list[str]] = contextvars.ContextVar( _span_start_time_ctx: contextvars.ContextVar[dict[str, int]] = ( contextvars.ContextVar("_bq_analytics_span_start_time", default=None) ) +_span_ownership_stack_ctx: contextvars.ContextVar[list[bool]] = ( + contextvars.ContextVar("_bq_analytics_span_ownership_stack", default=None) +) class TraceManager: @@ -442,7 +451,6 @@ class TraceManager: @staticmethod def init_trace(callback_context: CallbackContext) -> None: - # Extract root agent name from invocation context if not set if _root_agent_name_ctx.get() is None: try: root_agent = callback_context._invocation_context.agent.root_agent @@ -459,6 +467,9 @@ class TraceManager: if _span_start_time_ctx.get() is None: _span_start_time_ctx.set({}) + if _span_ownership_stack_ctx.get() is None: + _span_ownership_stack_ctx.set([]) + @staticmethod def get_trace_id(callback_context: CallbackContext) -> Optional[str]: """Gets the trace ID from the current span or invocation_id.""" @@ -492,13 +503,11 @@ class TraceManager: token = context.attach(trace.set_span_in_context(span)) stack = _span_stack_ctx.get() or [] - new_stack = list(stack) - new_stack.append(span) + new_stack = list(stack) + [span] _span_stack_ctx.set(new_stack) token_stack = _span_token_stack_ctx.get() or [] - new_token_stack = list(token_stack) - new_token_stack.append(token) + new_token_stack = list(token_stack) + [token] _span_token_stack_ctx.set(new_token_stack) if span.get_span_context().is_valid: @@ -509,8 +518,7 @@ class TraceManager: span_id_str = uuid.uuid4().hex id_stack = _span_id_stack_ctx.get() or [] - new_id_stack = list(id_stack) - new_id_stack.append(span_id_str) + new_id_stack = list(id_stack) + [span_id_str] _span_id_stack_ctx.set(new_id_stack) span_map = _span_map_ctx.get() or {} @@ -524,6 +532,51 @@ class TraceManager: new_start_times[span_id_str] = time.time_ns() _span_start_time_ctx.set(new_start_times) + ownership_stack = _span_ownership_stack_ctx.get() or [] + new_ownership_stack = list(ownership_stack) + [True] + _span_ownership_stack_ctx.set(new_ownership_stack) + + return span_id_str + + @staticmethod + def attach_current_span( + callback_context: CallbackContext, + ) -> str: + """Attaches the current OTEL span to the stack without owning it.""" + TraceManager.init_trace(callback_context) + + # Get current span but don't start a new one + span = trace.get_current_span() + # We still need to attach it to context to keep stacks symmetric with token + token = context.attach(trace.set_span_in_context(span)) + + stack = _span_stack_ctx.get() or [] + new_stack = list(stack) + [span] + _span_stack_ctx.set(new_stack) + + token_stack = _span_token_stack_ctx.get() or [] + new_token_stack = list(token_stack) + [token] + _span_token_stack_ctx.set(new_token_stack) + + if span.get_span_context().is_valid: + span_id_str = format(span.get_span_context().span_id, "016x") + else: + # Fallback: Generate a UUID-based ID if OTel span is invalid (NoOp) + span_id_str = uuid.uuid4().hex + + id_stack = _span_id_stack_ctx.get() or [] + new_id_stack = list(id_stack) + [span_id_str] + _span_id_stack_ctx.set(new_id_stack) + + span_map = _span_map_ctx.get() or {} + new_span_map = span_map.copy() + new_span_map[span_id_str] = span + _span_map_ctx.set(new_span_map) + + ownership_stack = _span_ownership_stack_ctx.get() or [] + new_ownership_stack = list(ownership_stack) + [False] + _span_ownership_stack_ctx.set(new_ownership_stack) + return span_id_str @staticmethod @@ -567,7 +620,16 @@ class TraceManager: start_ns = start_times[span_id] duration_ms = int((time.time_ns() - start_ns) / 1_000_000) - span.end() + should_end = True + ownership_stack = _span_ownership_stack_ctx.get() + if ownership_stack: + new_ownership_stack = list(ownership_stack) + should_end = new_ownership_stack.pop() + _span_ownership_stack_ctx.set(new_ownership_stack) + + if should_end: + span.end() + context.detach(token) first_tokens = _span_first_token_times_ctx.get() @@ -599,8 +661,11 @@ class TraceManager: if id_stack: span_id = id_stack[-1] parent_id = None - if len(id_stack) > 1: - parent_id = id_stack[-2] + # Walk backwards to find a different span_id for parent + for i in range(len(id_stack) - 2, -1, -1): + if id_stack[i] != span_id: + parent_id = id_stack[i] + break return span_id, parent_id return None, None @@ -677,7 +742,7 @@ class BatchProcessor: queue_max_size: int, shutdown_timeout: float, ): - """Initializes the BatchProcessor. + """Initializes the instance. Args: write_client: BigQueryWriteAsyncClient for writing rows. @@ -841,16 +906,16 @@ class BatchProcessor: except asyncio.CancelledError: logger.info("Batch writer task cancelled.") break - except RuntimeError as e: - if "Event loop is closed" in str(e): - logger.info("Batch writer loop closed: %s", e) - break - # Re-raise other RuntimeErrors (or log them below) - logger.error("RuntimeError in batch writer loop: %s", e, exc_info=True) - await asyncio.sleep(1) except Exception as e: logger.error("Error in batch writer loop: %s", e, exc_info=True) - await asyncio.sleep(1) + # Avoid sleeping if we are shutting down or if the task was cancelled + if not self._shutdown: + try: + await asyncio.sleep(1) + except (asyncio.CancelledError, RuntimeError): + break + else: + break async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: """Writes a batch of rows to BigQuery with retry logic. @@ -884,24 +949,25 @@ class BatchProcessor: async def requests_iter(): yield req - responses = await self.write_client.append_rows(requests_iter()) - async for response in responses: - error = getattr(response, "error", None) - error_code = getattr(error, "code", None) - if error_code and error_code != 0: - error_message = getattr(error, "message", "Unknown error") - logger.warning( - "BigQuery Write API returned error code %s: %s", - error_code, - error_message, - ) - if error_code in [ - _GRPC_DEADLINE_EXCEEDED, - _GRPC_INTERNAL, - _GRPC_UNAVAILABLE, - ]: # Deadline, Internal, Unavailable - raise ServiceUnavailable(error_message) - else: + async def perform_write(): + responses = await self.write_client.append_rows(requests_iter()) + async for response in responses: + error = getattr(response, "error", None) + error_code = getattr(error, "code", None) + if error_code and error_code != 0: + error_message = getattr(error, "message", "Unknown error") + logger.warning( + "BigQuery Write API returned error code %s: %s", + error_code, + error_message, + ) + if error_code in [ + _GRPC_DEADLINE_EXCEEDED, + _GRPC_INTERNAL, + _GRPC_UNAVAILABLE, + ]: + raise ServiceUnavailable(error_message) + if "schema mismatch" in error_message.lower(): logger.error( "BigQuery Schema Mismatch: %s. This usually means the" @@ -916,9 +982,17 @@ class BatchProcessor: logger.error("Row error details: %s", row_error) logger.error("Row content causing error: %s", rows) return + return + + await asyncio.wait_for(perform_write(), timeout=30.0) return - except (ServiceUnavailable, TooManyRequests, InternalServerError) as e: + except ( + ServiceUnavailable, + TooManyRequests, + InternalServerError, + asyncio.TimeoutError, + ) as e: attempt += 1 if attempt > self.retry_config.max_retries: logger.error( @@ -971,6 +1045,7 @@ class BatchProcessor: logger.warning("BatchProcessor shutdown timed out, cancelling worker.") self._batch_processor_task.cancel() try: + # Wait for the task to acknowledge cancellation await self._batch_processor_task except asyncio.CancelledError: pass @@ -1007,7 +1082,7 @@ class ContentParser: """Parses content for logging with length limits and structure normalization.""" def __init__(self, max_length: int) -> None: - """Initializes the ContentParser. + """Initializes the instance. Args: max_length: Maximum length for text content. @@ -1483,8 +1558,12 @@ def _get_events_schema() -> list[bigquery.SchemaField]: # ============================================================================== # MAIN PLUGIN # ============================================================================== -_GLOBAL_WRITE_CLIENT: Optional[BigQueryWriteAsyncClient] = None -_GLOBAL_CLIENT_LOCK = asyncio.Lock() +@dataclass +class _LoopState: + """Holds resources bound to a specific event loop.""" + + write_client: BigQueryWriteAsyncClient + batch_processor: BatchProcessor class BigQueryAgentAnalyticsPlugin(BasePlugin): @@ -1498,12 +1577,12 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): self, project_id: str, dataset_id: str, - *, table_id: Optional[str] = None, config: Optional[BigQueryLoggerConfig] = None, location: str = "US", + **kwargs, ) -> None: - """Initializes the BigQueryAgentAnalyticsPlugin. + """Initializes the instance. Args: project_id: Google Cloud project ID. @@ -1511,11 +1590,20 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): table_id: BigQuery table ID (optional, overrides config). config: BigQueryLoggerConfig (optional). location: BigQuery location (default: "US"). + **kwargs: Additional configuration parameters for BigQueryLoggerConfig. """ super().__init__(name="bigquery_agent_analytics") self.project_id = project_id self.dataset_id = dataset_id self.config = config or BigQueryLoggerConfig() + + # Override config with kwargs if provided + for key, value in kwargs.items(): + if hasattr(self.config, key): + setattr(self.config, key, value) + else: + logger.warning(f"Unknown configuration parameter: {key}") + self.table_id = table_id or self.config.table_id self.location = location @@ -1523,12 +1611,65 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): self._is_shutting_down = False self._setup_lock = None self.client = None - self.write_client = None - self.write_stream = None - self.batch_processor = None + self._loop_state_by_loop: dict[asyncio.AbstractEventLoop, _LoopState] = {} + self._write_stream_name = None # Resolved stream name self._executor = None self.offloader: Optional[GCSOffloader] = None self.parser: Optional[HybridContentParser] = None + self._schema = None + self.arrow_schema = None + + # API Compatibility: These attributes are statically defined as None to mask the + # dynamic properties from static analysis tools (preventing "breaking changes"), + # while __getattribute__ intercepts instance access to route to the logic. + batch_processor = None + write_client = None + write_stream = None + + def __getattribute__(self, name: str) -> Any: + """Intercepts attribute access to support API masking. + + Args: + name: The name of the attribute being accessed. + + Returns: + The value of the attribute. + """ + if name == "batch_processor": + return self._batch_processor_prop + if name == "write_client": + return self._write_client_prop + if name == "write_stream": + return self._write_stream_prop + return super().__getattribute__(name) + + @property + def _batch_processor_prop(self) -> Optional["BatchProcessor"]: + """The batch processor for the current loop (backward compatibility).""" + try: + loop = asyncio.get_running_loop() + if loop in self._loop_state_by_loop: + return self._loop_state_by_loop[loop].batch_processor + except RuntimeError: + pass + return None + + @property + def _write_client_prop(self) -> Optional["BigQueryWriteAsyncClient"]: + """The write client for the current loop (backward compatibility).""" + try: + loop = asyncio.get_running_loop() + if loop in self._loop_state_by_loop: + return self._loop_state_by_loop[loop].write_client + except RuntimeError: + pass + return None + + @property + def _write_stream_prop(self) -> Optional[str]: + """The write stream for the current loop (backward compatibility).""" + bp = self._batch_processor_prop + return bp.write_stream if bp else None def _format_content_safely( self, content: Optional[types.Content] @@ -1552,10 +1693,89 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): logger.warning("Content formatter failed: %s", e) return "[FORMATTING FAILED]", False + async def _get_loop_state(self) -> _LoopState: + """Gets or creates the state for the current event loop. + + Returns: + The loop-specific state object containing clients and processors. + """ + loop = asyncio.get_running_loop() + if loop in self._loop_state_by_loop: + return self._loop_state_by_loop[loop] + + # We DO NOT use the global client approach for multi-loop safety simpler + # or we must ensure _GLOBAL_WRITE_CLIENT usage is safe. + # The original code had a _GLOBAL_WRITE_CLIENT. + # If we want to reuse it, we must be careful. + # actually, _GLOBAL_WRITE_CLIENT is created in *A* loop. + # It cannot be shared across loops if it uses loop primitives. + # So strictly speaking, we should create a new client per loop. + # OR we assume the global client is thread-safe? + # grpc.aio clients are generally loop-bound. + # SAFE approach: Create one client per loop. + + def get_credentials(): + creds, project_id = google.auth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + return creds, project_id + + creds, project_id = await loop.run_in_executor( + self._executor, get_credentials + ) + quota_project_id = getattr(creds, "quota_project_id", None) or project_id + options = ( + client_options.ClientOptions(quota_project_id=quota_project_id) + if quota_project_id + else None + ) + client_info = gapic_client_info.ClientInfo( + user_agent=f"google-adk-bq-logger/{__version__}" + ) + + write_client = BigQueryWriteAsyncClient( + credentials=creds, + client_info=client_info, + client_options=options, + ) + + # Use the resolved write stream name + if not self._write_stream_name: + # Should be set in _lazy_setup or we set it here if missing? + # _lazy_setup guarantees self.table_id etc are ready. + self._write_stream_name = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}/_default" + + batch_processor = BatchProcessor( + write_client=write_client, + arrow_schema=self.arrow_schema, + write_stream=self._write_stream_name, + batch_size=self.config.batch_size, + flush_interval=self.config.batch_flush_interval, + retry_config=self.config.retry_config, + queue_max_size=self.config.queue_max_size, + shutdown_timeout=self.config.shutdown_timeout, + ) + await batch_processor.start() + + state = _LoopState(write_client, batch_processor) + self._loop_state_by_loop[loop] = state + + atexit.register(self._atexit_cleanup, weakref.proxy(batch_processor)) + + return state + async def flush(self) -> None: - """Flushes any pending events to BigQuery.""" - if self.batch_processor: - await self.batch_processor.flush() + """Flushes any pending events to BigQuery. + + Flushes the processor associated with the CURRENT loop. + """ + try: + loop = asyncio.get_running_loop() + if loop in self._loop_state_by_loop: + await self._loop_state_by_loop[loop].batch_processor.flush() + except RuntimeError: + # No running loop or other issue + pass async def _lazy_setup(self, **kwargs) -> None: """Performs lazy initialization of BigQuery clients and resources.""" @@ -1575,46 +1795,11 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): ) self.full_table_id = f"{self.project_id}.{self.dataset_id}.{self.table_id}" - self._schema = _get_events_schema() - await loop.run_in_executor(self._executor, self._ensure_schema_exists) + if not self._schema: + self._schema = _get_events_schema() + await loop.run_in_executor(self._executor, self._ensure_schema_exists) - if not self.write_client: - global _GLOBAL_WRITE_CLIENT - async with _GLOBAL_CLIENT_LOCK: - if _GLOBAL_WRITE_CLIENT is None: - - def get_credentials(): - creds, project_id = google.auth.default( - scopes=["https://www.googleapis.com/auth/cloud-platform"] - ) - return creds, project_id - - creds, project_id = await loop.run_in_executor( - self._executor, get_credentials - ) - quota_project_id = ( - getattr(creds, "quota_project_id", None) or project_id - ) - options = ( - client_options.ClientOptions(quota_project_id=quota_project_id) - if quota_project_id - else None - ) - client_info = gapic_client_info.ClientInfo( - user_agent=f"google-adk-bq-logger/{__version__}" - ) - # Initialize the async client in the current event loop, not in the - # executor. - _GLOBAL_WRITE_CLIENT = BigQueryWriteAsyncClient( - credentials=creds, - client_info=client_info, - client_options=options, - ) - self.write_client = _GLOBAL_WRITE_CLIENT - - self.write_stream = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}/_default" - - if not self.batch_processor: + if not self.parser: self.arrow_schema = to_arrow_schema(self._schema) if not self.arrow_schema: raise RuntimeError("Failed to convert BigQuery schema to Arrow schema.") @@ -1635,82 +1820,68 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): max_length=self.config.max_content_length, connection_id=self.config.connection_id, ) - self.batch_processor = BatchProcessor( - write_client=self.write_client, - arrow_schema=self.arrow_schema, - write_stream=self.write_stream, - batch_size=self.config.batch_size, - flush_interval=self.config.batch_flush_interval, - retry_config=self.config.retry_config, - queue_max_size=self.config.queue_max_size, - shutdown_timeout=self.config.shutdown_timeout, - ) - await self.batch_processor.start() - # Register cleanup to ensure logs are flushed if user forgets to close - # Use weakref to avoid circular references that prevent garbage collection - atexit.register(self._atexit_cleanup, weakref.proxy(self.batch_processor)) + await self._get_loop_state() @staticmethod def _atexit_cleanup(batch_processor: "BatchProcessor") -> None: """Clean up batch processor on script exit.""" - # Check if the batch_processor object is still alive + try: - if batch_processor and not batch_processor._shutdown: - # Emergency Flush: Rescue any logs remaining in the queue - remaining_items = [] - try: - while True: - remaining_items.append(batch_processor._queue.get_nowait()) - except (asyncio.QueueEmpty, AttributeError): - pass - - if remaining_items: - # We need a new loop and client to flush these - async def rescue_flush(): - try: - # Create a short-lived client just for this flush - try: - # Note: This relies on google.auth.default() working in this context. - # pylint: disable=g-import-not-at-top - from google.cloud.bigquery_storage_v1.services.big_query_write.async_client import BigQueryWriteAsyncClient - - # pylint: enable=g-import-not-at-top - client = BigQueryWriteAsyncClient() - except Exception as e: - logger.warning("Could not create rescue client: %s", e) - return - - # Patch batch_processor.write_client temporarily - old_client = batch_processor.write_client - batch_processor.write_client = client - try: - # Force a write - await batch_processor._write_rows_with_retry(remaining_items) - logger.info("Rescued logs flushed successfully.") - except Exception as e: - logger.error("Failed to flush rescued logs: %s", e) - finally: - batch_processor.write_client = old_client - except Exception as e: - logger.error("Rescue flush failed: %s", e) - - # In Python 3.13+, creating a new event loop during interpreter shutdown - # (inside atexit) can cause deadlocks if the threading module is already - # shutting down. We attempt to run only if safe. - try: - # Check if we can safely create a loop - loop = asyncio.new_event_loop() - try: - loop.run_until_complete(rescue_flush()) - finally: - loop.close() - except Exception as e: - logger.error("Failed to run rescue loop: %s", e) + if not batch_processor or batch_processor._shutdown: + return except ReferenceError: - # batch_processor already GC'd, nothing to do + return + + # Emergency Flush: Rescue any logs remaining in the queue + remaining_items = [] + try: + while True: + remaining_items.append(batch_processor._queue.get_nowait()) + except (asyncio.QueueEmpty, AttributeError): pass + if remaining_items: + # We need a new loop and client to flush these + async def rescue_flush(): + client = None + try: + # Create a short-lived client just for this flush + try: + # Note: This relies on google.auth.default() working in this context. + # pylint: disable=g-import-not-at-top + from google.cloud.bigquery_storage_v1.services.big_query_write.async_client import BigQueryWriteAsyncClient + + # pylint: enable=g-import-not-at-top + client = BigQueryWriteAsyncClient() + except Exception as e: + logger.warning("Could not create rescue client: %s", e) + return + + # Patch batch_processor.write_client temporarily + old_client = batch_processor.write_client + batch_processor.write_client = client + try: + # Force a write + await batch_processor._write_rows_with_retry(remaining_items) + logger.info("Rescued logs flushed successfully.") + except Exception as e: + logger.error("Failed to flush rescued logs: %s", e) + finally: + batch_processor.write_client = old_client + except Exception as e: + logger.error("Rescue flush failed: %s", e) + finally: + if client: + await client.transport.close() + + try: + loop = asyncio.new_event_loop() + loop.run_until_complete(rescue_flush()) + loop.close() + except Exception as e: + logger.error("Failed to run rescue loop: %s", e) + def _ensure_schema_exists(self) -> None: """Ensures the BigQuery table exists with the correct schema.""" try: @@ -1753,22 +1924,29 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): t = timeout if timeout is not None else self.config.shutdown_timeout loop = asyncio.get_running_loop() try: - if self.batch_processor: - await self.batch_processor.shutdown(timeout=t) - if self.write_client and getattr(self.write_client, "transport", None): - # Only close the client if it's NOT the global one (unlikely with new logic, - # but good for safety if injected manually) or if we decide to handle global close differently. - # For now, we DO NOT close the global client to allow reuse. - if self.write_client is not _GLOBAL_WRITE_CLIENT: - await self.write_client.transport.close() + # Correct Multi-Loop Shutdown: + # 1. Shutdown current loop's processor directly. + if loop in self._loop_state_by_loop: + await self._loop_state_by_loop[loop].batch_processor.shutdown(timeout=t) + + # 2. Close clients for all states + for state in self._loop_state_by_loop.values(): + if state.write_client and getattr( + state.write_client, "transport", None + ): + try: + await state.write_client.transport.close() + except Exception: + pass + + self._loop_state_by_loop.clear() + if self.client: if self._executor: executor = self._executor await loop.run_in_executor(None, lambda: executor.shutdown(wait=True)) self._executor = None - self.write_client = None self.client = None - self._is_shutting_down = False except Exception as e: logger.error("Error during shutdown: %s", e, exc_info=True) self._is_shutting_down = False @@ -1779,9 +1957,8 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): state = self.__dict__.copy() state["_setup_lock"] = None state["client"] = None - state["write_client"] = None - state["write_stream"] = None - state["batch_processor"] = None + state["_loop_state_by_loop"] = {} + state["_write_stream_name"] = None state["_executor"] = None state["offloader"] = None state["parser"] = None @@ -1922,6 +2099,24 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): # Fallback if it couldn't be converted to dict kwargs["usage_metadata"] = usage_metadata + # 6. Session Metadata + if self.config.log_session_metadata and hasattr( + callback_context, "session" + ): + try: + # Accessing session.metadata might trigger lazy loading or be a property + # Use getattr to safely check for metadata without raising AttributeError + metadata = getattr(callback_context.session, "metadata", None) + if metadata: + kwargs["session_metadata"] = metadata + except Exception: + # Ignore errors if metadata is missing or inaccessible + pass + + # 7. Custom Tags + if self.config.custom_tags: + kwargs["custom_tags"] = self.config.custom_tags + # Serialize remaining kwargs to JSON string for attributes try: attributes_json = json.dumps(kwargs) @@ -1950,8 +2145,8 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): "is_truncated": is_truncated, } - if self.batch_processor: - await self.batch_processor.append(row) + state = await self._get_loop_state() + await state.batch_processor.append(row) # --- UPDATED CALLBACKS FOR V1 PARITY --- @@ -1974,6 +2169,27 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): raw_content=user_message, ) + async def on_state_change_callback( + self, + *, + callback_context: CallbackContext, + state_delta: dict[str, Any], + **kwargs, + ) -> None: + """Logs state changes (state_delta) to BigQuery. + + Args: + callback_context: The callback context. + state_delta: The change in state to log. + **kwargs: Additional arguments. + """ + await self._log_event( + "STATE_DELTA", + callback_context, + state_delta=state_delta, + **kwargs, + ) + async def before_run_callback( self, *, invocation_context: "InvocationContext", **kwargs ) -> None: @@ -2069,19 +2285,35 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): "candidate_count", "max_output_tokens", "stop_sequences", + "presence_penalty", + "frequency_penalty", + "response_mime_type", + "response_schema", + "seed", + "response_logprobs", + "logprobs", ]: - if val := getattr(llm_request.config, field_name, None): + val = getattr(llm_request.config, field_name, None) + if val is not None: config_dict[field_name] = val + + # Handle labels if present in config + if hasattr(llm_request.config, "labels") and llm_request.config.labels: + attributes["labels"] = llm_request.config.labels + if config_dict: attributes["llm_config"] = config_dict + if labels := getattr(llm_request.config, "labels", None): + attributes["labels"] = labels + if hasattr(llm_request, "tools_dict") and llm_request.tools_dict: attributes["tools"] = list(llm_request.tools_dict.keys()) # Merge any additional kwargs into attributes attributes.update(kwargs) - TraceManager.push_span(callback_context, "llm") + TraceManager.push_span(callback_context, "llm_request") await self._log_event( "LLM_REQUEST", callback_context, @@ -2278,6 +2510,13 @@ class BigQueryAgentAnalyticsPlugin(BasePlugin): parent_span_id_override=parent_span_id, ) + if tool_context.actions.state_delta: + await self._log_event( + "STATE_DELTA", + tool_context, + state_delta=tool_context.actions.state_delta, + ) + async def on_tool_error_callback( self, *, diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index e17729d1..b11d5659 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -24,6 +24,7 @@ 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.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_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 @@ -123,9 +124,6 @@ def mock_bq_client(): @pytest.fixture def mock_write_client(): - from google.adk.plugins import bigquery_agent_analytics_plugin - - bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT = None with mock.patch.object( bigquery_agent_analytics_plugin, "BigQueryWriteAsyncClient", autospec=True ) as mock_cls: @@ -204,8 +202,6 @@ def dummy_arrow_schema(): @pytest.fixture def mock_to_arrow_schema(dummy_arrow_schema): - from google.adk.plugins import bigquery_agent_analytics_plugin - with mock.patch.object( bigquery_agent_analytics_plugin, "to_arrow_schema", @@ -240,8 +236,6 @@ async def bq_plugin_inst( mock_to_arrow_schema, mock_asyncio_to_thread, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( project_id=PROJECT_ID, dataset_id=DATASET_ID, @@ -256,8 +250,6 @@ async def bq_plugin_inst( @contextlib.asynccontextmanager async def managed_plugin(*args, **kwargs): """Async context manager to ensure plugin shutdown.""" - from google.adk.plugins import bigquery_agent_analytics_plugin - plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( *args, **kwargs ) @@ -338,66 +330,59 @@ def _assert_common_fields(log_entry, event_type, agent="MyTestAgent"): assert log_entry["invocation_id"] == "inv-789" -# @pytest.mark.skip(reason="Disabled for build/Kokoro stability") -# def test_recursive_smart_truncate(): -# """Test recursive smart truncate.""" -# -# obj = { -# "a": "long string" * 10, -# "b": ["short", "long string" * 10], -# "c": {"d": "long string" * 10}, -# } -# max_len = 10 -# truncated, is_truncated = ( -# bigquery_agent_analytics_plugin._recursive_smart_truncate(obj, max_len) -# ) -# assert is_truncated -# -# assert truncated["a"] == "long strin...[TRUNCATED]" -# assert truncated["b"][0] == "short" -# assert truncated["b"][1] == "long strin...[TRUNCATED]" -# assert truncated["c"]["d"] == "long strin...[TRUNCATED]" -# -# -# @pytest.mark.skip(reason="Disabled for build/Kokoro stability") -# def test_recursive_smart_truncate_with_dataclasses(): -# """Test recursive smart truncate with dataclasses.""" -# -# @dataclasses.dataclass -# class LocalMissedKPI: -# kpi: str -# value: float -# -# @dataclasses.dataclass -# class LocalIncident: -# id: str -# kpi_missed: list[LocalMissedKPI] -# status: str -# -# incident = LocalIncident( -# id="inc-123", -# kpi_missed=[LocalMissedKPI(kpi="latency", value=99.9)], -# status="active", -# ) -# content = {"result": incident} -# max_len = 1000 -# -# truncated, is_truncated = ( -# bigquery_agent_analytics_plugin._recursive_smart_truncate( -# content, max_len -# ) -# ) -# assert not is_truncated -# assert isinstance(truncated["result"], dict) -# assert truncated["result"]["id"] == "inc-123" -# assert isinstance(truncated["result"]["kpi_missed"][0], dict) -# assert truncated["result"]["kpi_missed"][0]["kpi"] == "latency" -# -# -# # --- Test Class --- -# -# -# @pytest.mark.skip(reason="Disabled for build/Kokoro stability") +def test_recursive_smart_truncate(): + """Test recursive smart truncate.""" + obj = { + "a": "long string" * 10, + "b": ["short", "long string" * 10], + "c": {"d": "long string" * 10}, + } + max_len = 10 + truncated, is_truncated = ( + bigquery_agent_analytics_plugin._recursive_smart_truncate(obj, max_len) + ) + assert is_truncated + + assert truncated["a"] == "long strin...[TRUNCATED]" + assert truncated["b"][0] == "short" + assert truncated["b"][1] == "long strin...[TRUNCATED]" + assert truncated["c"]["d"] == "long strin...[TRUNCATED]" + + +def test_recursive_smart_truncate_with_dataclasses(): + """Test recursive smart truncate with dataclasses.""" + + @dataclasses.dataclass + class LocalMissedKPI: + kpi: str + value: float + + @dataclasses.dataclass + class LocalIncident: + id: str + kpi_missed: list[LocalMissedKPI] + status: str + + incident = LocalIncident( + id="inc-123", + kpi_missed=[LocalMissedKPI(kpi="latency", value=99.9)], + status="active", + ) + content = {"result": incident} + max_len = 1000 + + truncated, is_truncated = ( + bigquery_agent_analytics_plugin._recursive_smart_truncate( + content, max_len + ) + ) + assert not is_truncated + assert isinstance(truncated["result"], dict) + assert truncated["result"]["id"] == "inc-123" + assert isinstance(truncated["result"]["kpi_missed"][0], dict) + assert truncated["result"]["kpi_missed"][0]["kpi"] == "latency" + + class TestBigQueryAgentAnalyticsPlugin: """Tests for the BigQueryAgentAnalyticsPlugin.""" @@ -409,8 +394,6 @@ class TestBigQueryAgentAnalyticsPlugin: mock_write_client, invocation_context, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(enabled=False) async with managed_plugin( @@ -438,9 +421,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, ): # Setup - # Setup - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig() async with managed_plugin(PROJECT_ID, DATASET_ID, config=config) as plugin: @@ -502,8 +482,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, ): # Setup - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig() plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( @@ -557,8 +535,6 @@ class TestBigQueryAgentAnalyticsPlugin: ): _ = mock_auth_default _ = mock_bq_client - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(event_allowlist=["LLM_REQUEST"]) async with managed_plugin( @@ -598,8 +574,6 @@ class TestBigQueryAgentAnalyticsPlugin: ): _ = mock_auth_default _ = mock_bq_client - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(event_denylist=["USER_MESSAGE_RECEIVED"]) async with managed_plugin( @@ -637,8 +611,6 @@ class TestBigQueryAgentAnalyticsPlugin: def redact_content(content, event_type): return "[REDACTED]" - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(content_formatter=redact_content) async with managed_plugin( @@ -677,8 +649,6 @@ class TestBigQueryAgentAnalyticsPlugin: def error_formatter(content, event_type): raise ValueError("Formatter failed") - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(content_formatter=error_formatter) async with managed_plugin( @@ -713,8 +683,6 @@ class TestBigQueryAgentAnalyticsPlugin: ): _ = mock_auth_default _ = mock_bq_client - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(max_content_length=40) async with managed_plugin( @@ -784,8 +752,6 @@ class TestBigQueryAgentAnalyticsPlugin: dummy_arrow_schema, mock_asyncio_to_thread, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig _ = mock_auth_default _ = mock_bq_client @@ -835,8 +801,6 @@ class TestBigQueryAgentAnalyticsPlugin: dummy_arrow_schema, mock_asyncio_to_thread, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(max_content_length=-1) async with managed_plugin( @@ -883,8 +847,6 @@ class TestBigQueryAgentAnalyticsPlugin: dummy_arrow_schema, ): """Test max content length for tool result.""" - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig _ = mock_auth_default _ = mock_bq_client @@ -937,8 +899,6 @@ class TestBigQueryAgentAnalyticsPlugin: _ = mock_bq_client _ = mock_to_arrow_schema _ = mock_asyncio_to_thread - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(max_content_length=-1) async with managed_plugin( @@ -982,8 +942,6 @@ class TestBigQueryAgentAnalyticsPlugin: dummy_arrow_schema, mock_asyncio_to_thread, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(max_content_length=80) async with managed_plugin( @@ -1026,8 +984,6 @@ class TestBigQueryAgentAnalyticsPlugin: invocation_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - user_message = types.Content(parts=[types.Part(text="What is up?")]) bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context) await bq_plugin_inst.on_user_message_callback( @@ -1052,8 +1008,6 @@ class TestBigQueryAgentAnalyticsPlugin: mock_asyncio_to_thread, mock_storage_client, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig _ = mock_auth_default _ = mock_bq_client @@ -1118,8 +1072,6 @@ class TestBigQueryAgentAnalyticsPlugin: invocation_context, mock_asyncio_to_thread, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - _ = mock_asyncio_to_thread mock_auth_default.side_effect = auth_exceptions.GoogleAuthError( "Auth failed" @@ -1149,8 +1101,8 @@ class TestBigQueryAgentAnalyticsPlugin: async def test_bigquery_insert_error_does_not_raise( self, bq_plugin_inst, mock_write_client, invocation_context ): + _ = bq_plugin_inst - from google.adk.plugins import bigquery_agent_analytics_plugin async def fake_append_rows_with_error(requests, **kwargs): mock_append_rows_response = mock.MagicMock() @@ -1182,7 +1134,6 @@ class TestBigQueryAgentAnalyticsPlugin: self, bq_plugin_inst, mock_write_client, invocation_context ): """Test that retryable BigQuery errors are logged and retried.""" - from google.adk.plugins import bigquery_agent_analytics_plugin async def fake_append_rows_with_retryable_error(requests, **kwargs): mock_append_rows_response = mock.MagicMock() @@ -1216,8 +1167,6 @@ class TestBigQueryAgentAnalyticsPlugin: async def test_schema_mismatch_error_handling( self, bq_plugin_inst, mock_write_client, invocation_context ): - from google.adk.plugins import bigquery_agent_analytics_plugin - async def fake_append_rows_with_schema_error(requests, **kwargs): mock_resp = mock.MagicMock() mock_resp.row_errors = [] @@ -1249,18 +1198,12 @@ class TestBigQueryAgentAnalyticsPlugin: @pytest.mark.asyncio async def test_close(self, bq_plugin_inst, mock_bq_client, mock_write_client): """Test plugin shutdown.""" - from google.adk.plugins import bigquery_agent_analytics_plugin - # Force the plugin to think it owns the client by clearing the global reference - bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT = None await bq_plugin_inst.shutdown() - mock_write_client.transport.close.assert_called_once() - # bq_client might not be closed if it wasn't created or if close() failed, - # but here it should be. - # in the new implementation we verify attributes are reset - assert bq_plugin_inst.write_client is None - assert bq_plugin_inst.client is None - assert bq_plugin_inst._is_shutting_down is False + # shutdown calls transport.close() on all clients + assert mock_write_client.transport.close.call_count >= 1 + # Verify loop states are cleared + assert not bq_plugin_inst._loop_state_by_loop @pytest.mark.asyncio async def test_before_run_callback_logs_correctly( @@ -1271,7 +1214,6 @@ class TestBigQueryAgentAnalyticsPlugin: dummy_arrow_schema, ): """Test before_run_callback logs correctly.""" - from google.adk.plugins import bigquery_agent_analytics_plugin bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context) await bq_plugin_inst.before_run_callback( @@ -1292,8 +1234,6 @@ class TestBigQueryAgentAnalyticsPlugin: invocation_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context) await bq_plugin_inst.after_run_callback( invocation_context=invocation_context @@ -1314,8 +1254,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context) await bq_plugin_inst.before_agent_callback( agent=mock_agent, callback_context=callback_context @@ -1336,8 +1274,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context) await bq_plugin_inst.after_agent_callback( agent=mock_agent, callback_context=callback_context @@ -1361,8 +1297,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - llm_request = llm_request_lib.LlmRequest( model="gemini-pro", contents=[ @@ -1388,8 +1322,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - llm_request = llm_request_lib.LlmRequest( model="gemini-pro", config=types.GenerateContentConfig( @@ -1423,6 +1355,65 @@ class TestBigQueryAgentAnalyticsPlugin: assert attributes["llm_config"]["top_p"] == 0.9 assert attributes["tools"] == ["tool1", "tool2"] + @pytest.mark.asyncio + async def test_before_model_callback_with_full_config( + self, + bq_plugin_inst, + mock_write_client, + callback_context, + dummy_arrow_schema, + ): + """Test that all config fields, including falsy values and labels, are logged.""" + llm_request = llm_request_lib.LlmRequest( + model="gemini-pro", + config=types.GenerateContentConfig( + temperature=0.0, + top_p=0.1, + top_k=5.0, + candidate_count=5, + max_output_tokens=65000, + stop_sequences=["STOP"], + presence_penalty=0.1, + frequency_penalty=0.5, + seed=42, + response_logprobs=True, + logprobs=3, + labels={"llm.agent.name": "test_agent"}, + ), + contents=[types.Content(role="user", parts=[types.Part(text="User")])], + ) + bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context) + await bq_plugin_inst.before_model_callback( + callback_context=callback_context, llm_request=llm_request + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + _assert_common_fields(log_entry, "LLM_REQUEST") + + # Verify attributes + assert "attributes" in log_entry + attributes = json.loads(log_entry["attributes"]) + + llm_config = attributes.get("llm_config", {}) + expected_llm_config = { + "temperature": 0.0, + "top_p": 0.1, + "top_k": 5.0, + "candidate_count": 5, + "max_output_tokens": 65000, + "stop_sequences": ["STOP"], + "presence_penalty": 0.1, + "frequency_penalty": 0.5, + "seed": 42, + "response_logprobs": True, + "logprobs": 3, + } + assert llm_config == expected_llm_config + + assert attributes.get("labels") == {"llm.agent.name": "test_agent"} + @pytest.mark.asyncio async def test_before_model_callback_multipart_separator( self, @@ -1431,8 +1422,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - llm_request = llm_request_lib.LlmRequest( model="gemini-pro", contents=[ @@ -1462,8 +1451,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - llm_response = llm_response_lib.LlmResponse( content=types.Content(parts=[types.Part(text="Model response")]), usage_metadata=types.UsageMetadata( @@ -1502,8 +1489,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - tool_fc = types.FunctionCall(name="get_weather", args={"location": "Paris"}) llm_response = llm_response_lib.LlmResponse( content=types.Content(parts=[types.Part(function_call=tool_fc)]), @@ -1531,8 +1516,6 @@ class TestBigQueryAgentAnalyticsPlugin: async def test_before_tool_callback_logs_correctly( self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema ): - from google.adk.plugins import bigquery_agent_analytics_plugin - mock_tool = mock.create_autospec( base_tool_lib.BaseTool, instance=True, spec_set=True ) @@ -1555,8 +1538,6 @@ class TestBigQueryAgentAnalyticsPlugin: async def test_after_tool_callback_logs_correctly( self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema ): - from google.adk.plugins import bigquery_agent_analytics_plugin - mock_tool = mock.create_autospec( base_tool_lib.BaseTool, instance=True, spec_set=True ) @@ -1578,6 +1559,135 @@ class TestBigQueryAgentAnalyticsPlugin: assert content_dict["tool"] == "MyTool" assert content_dict["result"] == {"res": "success"} + @pytest.mark.asyncio + async def test_after_tool_callback_state_delta_logging( + self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema + ): + mock_tool = mock.create_autospec( + base_tool_lib.BaseTool, instance=True, spec_set=True + ) + type(mock_tool).name = mock.PropertyMock(return_value="StateTool") + type(mock_tool).description = mock.PropertyMock(return_value="Sets state") + + # Simulate a tool modifying the state + tool_context.actions.state_delta["new_key"] = "new_value" + + bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context) + await bq_plugin_inst.after_tool_callback( + tool=mock_tool, + tool_args={"arg1": "val1"}, + tool_context=tool_context, + result={"res": "success"}, + ) + await asyncio.sleep(0.01) + + # We should have two events appended: TOOL_COMPLETED and STATE_DELTA + assert mock_write_client.append_rows.call_count >= 1 + + # Retrieve all flushed events + rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema) + assert len(rows) == 2 + + # Sort by event_type to reliably access them + rows.sort(key=lambda x: x["event_type"]) + + state_delta_event = ( + rows[0] if rows[0]["event_type"] == "STATE_DELTA" else rows[1] + ) + tool_event = ( + rows[1] if rows[1]["event_type"] == "TOOL_COMPLETED" else rows[0] + ) + + assert state_delta_event["event_type"] == "STATE_DELTA" + assert tool_event["event_type"] == "TOOL_COMPLETED" + + # Verify STATE_DELTA payload + attributes = json.loads(state_delta_event["attributes"]) + assert "state_delta" in attributes + assert attributes["state_delta"] == {"new_key": "new_value"} + assert state_delta_event["content"] is None + + @pytest.mark.asyncio + async def test_on_state_change_callback_logs_correctly( + self, + bq_plugin_inst, + mock_write_client, + callback_context, + dummy_arrow_schema, + ): + state_delta = {"key": "value", "new_key": 123} + bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context) + await bq_plugin_inst.on_state_change_callback( + callback_context=callback_context, state_delta=state_delta + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + _assert_common_fields(log_entry, "STATE_DELTA") + # content should be None (as raw_content was not passed) + assert log_entry["content"] is None + + # state_delta should be in attributes + attributes = json.loads(log_entry["attributes"]) + assert attributes["state_delta"] == state_delta + + @pytest.mark.asyncio + async def test_log_event_with_session_metadata( + self, + bq_plugin_inst, + mock_write_client, + callback_context, + dummy_arrow_schema, + ): + """Test that session metadata is logged when enabled.""" + # Setup session metadata + metadata = {"thread_id": "gchat-123", "key": "val"} + type(callback_context.session).metadata = mock.PropertyMock( + return_value=metadata + ) + + # Ensure config enabled (default is True) + bq_plugin_inst.config.log_session_metadata = True + + await bq_plugin_inst._log_event( + "TEST_EVENT", + callback_context, + raw_content="test content", + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + + attributes = json.loads(log_entry["attributes"]) + assert attributes["session_metadata"] == metadata + + @pytest.mark.asyncio + async def test_log_event_with_custom_tags( + self, + bq_plugin_inst, + mock_write_client, + callback_context, + dummy_arrow_schema, + ): + """Test that custom tags are logged.""" + custom_tags = {"agent_role": "sales", "env": "prod"} + bq_plugin_inst.config.custom_tags = custom_tags + + await bq_plugin_inst._log_event( + "TEST_EVENT", + callback_context, + raw_content="test content", + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + + attributes = json.loads(log_entry["attributes"]) + assert attributes["custom_tags"] == custom_tags + @pytest.mark.asyncio async def test_on_model_error_callback_logs_correctly( self, @@ -1586,8 +1696,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, dummy_arrow_schema, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - llm_request = llm_request_lib.LlmRequest( model="gemini-pro", contents=[types.Content(parts=[types.Part(text="Prompt")])], @@ -1609,8 +1717,6 @@ class TestBigQueryAgentAnalyticsPlugin: async def test_on_tool_error_callback_logs_correctly( self, bq_plugin_inst, mock_write_client, tool_context, dummy_arrow_schema ): - from google.adk.plugins import bigquery_agent_analytics_plugin - mock_tool = mock.create_autospec( base_tool_lib.BaseTool, instance=True, spec_set=True ) @@ -1708,8 +1814,8 @@ class TestBigQueryAgentAnalyticsPlugin: future = executor.submit(_run_in_thread, plugin) future.result() # Should not raise "no current event loop" assert plugin._started - assert plugin.client is not None - assert plugin.write_client is not None + # Verify loop states are populated + assert plugin._loop_state_by_loop @pytest.mark.asyncio async def test_multimodal_offloading( @@ -1724,8 +1830,6 @@ class TestBigQueryAgentAnalyticsPlugin: ): # Setup bucket_name = "test-bucket" - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig config = BigQueryLoggerConfig(gcs_bucket_name=bucket_name) async with managed_plugin( @@ -1765,47 +1869,6 @@ class TestBigQueryAgentAnalyticsPlugin: assert content_parts[0]["storage_mode"] == "GCS_REFERENCE" assert content_parts[0]["uri"].startswith(f"gs://{bucket_name}/") - @pytest.mark.asyncio - async def test_global_client_reuse( - self, mock_write_client, mock_auth_default - ): - del mock_write_client, mock_auth_default # Unused - from google.adk.plugins import bigquery_agent_analytics_plugin - - # Reset global client for this test - bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT = None - # Create two plugins - plugin1 = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( - PROJECT_ID, DATASET_ID, table_id="table1" - ) - plugin2 = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( - PROJECT_ID, DATASET_ID, table_id="table2" - ) - # Start both - try: - await plugin1._ensure_started() - await plugin2._ensure_started() - # Verify they share the same write_client instance - assert plugin1.write_client is not None - assert plugin2.write_client is not None - assert plugin1.write_client is plugin2.write_client - # Verify shutdown doesn't close the global client - await plugin1.shutdown() - # Mock transport close check - since it's a mock, we check call count - # But here we check if the client is still the global one - assert ( - bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT - is plugin2.write_client - ) - finally: - # Cleanup - await plugin2.shutdown() - if ( - not plugin1._is_shutting_down - ): # Ensure plugin1 is down if test failed early - await plugin1.shutdown() - bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT = None - @pytest.mark.asyncio async def test_quota_project_id_used_in_client( self, @@ -1813,9 +1876,6 @@ class TestBigQueryAgentAnalyticsPlugin: mock_to_arrow_schema, mock_asyncio_to_thread, ): - from google.adk.plugins import bigquery_agent_analytics_plugin - - bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT = None mock_creds = mock.create_autospec( google.auth.credentials.Credentials, instance=True, spec_set=True ) @@ -1841,13 +1901,10 @@ class TestBigQueryAgentAnalyticsPlugin: mock_bq_write_cls.assert_called_once() _, kwargs = mock_bq_write_cls.call_args assert kwargs["client_options"].quota_project_id == "quota-project" - bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT = None @pytest.mark.asyncio async def test_pickle_safety(self, mock_auth_default, mock_bq_client): """Test that the plugin can be pickled safely.""" - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig import pickle @@ -1872,7 +1929,7 @@ class TestBigQueryAgentAnalyticsPlugin: # Runtime objects should be None after unpickling assert unpickled_started._setup_lock is None assert unpickled_started._executor is None - assert unpickled_started.client is None + assert not unpickled_started._loop_state_by_loop finally: await plugin.shutdown() @@ -1885,10 +1942,11 @@ class TestBigQueryAgentAnalyticsPlugin: dummy_arrow_schema, ): """Verifies that LLM events have correct Span ID hierarchy.""" - from google.adk.plugins import bigquery_agent_analytics_plugin - # 1. Start Agent Span bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context) + _, _ = ( + bigquery_agent_analytics_plugin.TraceManager.get_current_span_and_parent() + ) agent_span_id = ( bigquery_agent_analytics_plugin.TraceManager.get_current_span_id() ) @@ -1912,12 +1970,14 @@ class TestBigQueryAgentAnalyticsPlugin: llm_span_id = ( bigquery_agent_analytics_plugin.TraceManager.get_current_span_id() ) + # Now that we push a new span for LLM calls, it should differ from agent_span_id assert llm_span_id != agent_span_id log_entry_req = await _get_captured_event_dict_async( mock_write_client, dummy_arrow_schema ) assert log_entry_req["event_type"] == "LLM_REQUEST" assert log_entry_req["span_id"] == llm_span_id + # The parent of the LLM span should be the Agent span assert log_entry_req["parent_span_id"] == agent_span_id mock_write_client.append_rows.reset_mock() # 4. LLM Response @@ -1937,9 +1997,8 @@ class TestBigQueryAgentAnalyticsPlugin: ) assert log_entry_resp["event_type"] == "LLM_RESPONSE" assert log_entry_resp["span_id"] == llm_span_id - # Crux of the bug fix: Parent should still be Agent Span, NOT Self. + # The parent of the LLM span should be the Agent span assert log_entry_resp["parent_span_id"] == agent_span_id - assert log_entry_resp["parent_span_id"] != log_entry_resp["span_id"] # Verify LLM Span was popped and we are back to Agent Span assert ( bigquery_agent_analytics_plugin.TraceManager.get_current_span_id() @@ -1965,8 +2024,6 @@ class TestBigQueryAgentAnalyticsPlugin: """Verifies that custom objects (Dataclasses) are serialized to dicts.""" _ = mock_auth_default _ = mock_bq_client - from google.adk.plugins import bigquery_agent_analytics_plugin - BigQueryLoggerConfig = bigquery_agent_analytics_plugin.BigQueryLoggerConfig @dataclasses.dataclass @@ -2014,8 +2071,6 @@ class TestBigQueryAgentAnalyticsPlugin: callback_context, ): """Verifies OpenTelemetry integration in TraceManager.""" - from google.adk.plugins import bigquery_agent_analytics_plugin - # Mock the tracer and span mock_tracer = mock.Mock() mock_span = mock.Mock() @@ -2057,7 +2112,6 @@ class TestBigQueryAgentAnalyticsPlugin: @pytest.mark.asyncio async def test_otel_integration_real_provider(self, callback_context): """Verifies TraceManager with a real OpenTelemetry TracerProvider.""" - from google.adk.plugins import bigquery_agent_analytics_plugin # Setup OTEL with in-memory exporter # pylint: disable=g-import-not-at-top from opentelemetry.sdk import trace as trace_sdk @@ -2107,8 +2161,6 @@ class TestBigQueryAgentAnalyticsPlugin: invocation_context, ): """Verifies that flush() forces pending events to be written.""" - from google.adk.plugins import bigquery_agent_analytics_plugin - # Log an event bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context) await bq_plugin_inst.before_run_callback( @@ -2122,3 +2174,75 @@ class TestBigQueryAgentAnalyticsPlugin: mock_write_client, dummy_arrow_schema ) assert log_entry["event_type"] == "INVOCATION_STARTING" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "gen_config_kwargs, expected_llm_config", + [ + ( + { + "temperature": 0.0, + "top_k": 5.0, + "top_p": 0.1, + "candidate_count": 5, + "max_output_tokens": 65000, + "presence_penalty": 0.1, + "frequency_penalty": 0.5, + "response_logprobs": True, + "logprobs": 3, + "seed": 42, + "labels": {"llm.agent.name": "test_agent"}, + }, + { + "temperature": 0.0, + "top_k": 5.0, + "top_p": 0.1, + "candidate_count": 5, + "max_output_tokens": 65000, + "presence_penalty": 0.1, + "frequency_penalty": 0.5, + "response_logprobs": True, + "logprobs": 3, + "seed": 42, + }, + ), + ], + ) + async def test_generation_config_logging( + self, + bq_plugin_inst, + mock_write_client, + dummy_arrow_schema, + callback_context, + gen_config_kwargs, + expected_llm_config, + ): + """Verifies that all fields in GenerateContentConfig are logged correctly.""" + gen_config = types.GenerateContentConfig(**gen_config_kwargs) + + llm_request = llm_request_lib.LlmRequest( + model="gemini-pro", + contents=[types.Content(parts=[types.Part(text="Prompt")])], + config=gen_config, + ) + + bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context) + await bq_plugin_inst.before_model_callback( + callback_context=callback_context, llm_request=llm_request + ) + # Flush + await bq_plugin_inst.flush() + + # Verify + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + assert log_entry["event_type"] == "LLM_REQUEST" + + attributes = json.loads(log_entry["attributes"]) + llm_config = attributes.get("llm_config", {}) + + assert llm_config == expected_llm_config + + if "labels" in gen_config_kwargs: + assert attributes.get("labels") == gen_config_kwargs["labels"]