encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity
Files
2367901ec54824984d9f0ef9aee711a6edbcb990
adk-python/src/google/adk/plugins/bigquery_agent_analytics_plugin.py
T
George WealeandCopybara-Service 2367901ec5 chore: Upgrade to headers to 2026 
Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 858763407
2026-01-20 14:50:09 -08:00

2101 lines
69 KiB
Python
Raw Blame History

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import asyncio
import atexit
from concurrent.futures import ThreadPoolExecutor
import contextvars
import dataclasses
from dataclasses import dataclass
from dataclasses import field
from datetime import datetime
from datetime import timezone
import functools
import json
import logging
import mimetypes
import random
import time
from types import MappingProxyType
from typing import Any
from typing import Callable
from typing import Optional
from typing import TYPE_CHECKING
import uuid
import weakref
from google.api_core import client_options
from google.api_core.exceptions import InternalServerError
from google.api_core.exceptions import ServiceUnavailable
from google.api_core.exceptions import TooManyRequests
from google.api_core.gapic_v1 import client_info as gapic_client_info
import google.auth
from google.cloud import bigquery
from google.cloud import exceptions as cloud_exceptions
from google.cloud import storage
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
from google.genai import types
import pyarrow as pa
from ..agents.callback_context import CallbackContext
from ..models.llm_request import LlmRequest
from ..models.llm_response import LlmResponse
from ..tools.base_tool import BaseTool
from ..tools.tool_context import ToolContext
from ..version import __version__
from .base_plugin import BasePlugin
if TYPE_CHECKING:
from ..agents.invocation_context import InvocationContext
logger: logging.Logger = logging.getLogger("google_adk." + __name__)
# gRPC Error Codes
_GRPC_DEADLINE_EXCEEDED = 4
_GRPC_INTERNAL = 13
_GRPC_UNAVAILABLE = 14
# --- Helper Formatters ---
def _format_content(
content: Optional[types.Content], *, max_len: int = 5000
) -> tuple[str, bool]:
"""Formats an Event content for logging.
Args:
content: The content to format.
max_len: Maximum length for text parts.
Returns:
A tuple of (formatted_string, is_truncated).
"""
if content is None or not content.parts:
return "None", False
parts = []
truncated = False
for p in content.parts:
if p.text:
if max_len != -1 and len(p.text) > max_len:
parts.append(f"text: '{p.text[:max_len]}...'")
truncated = True
else:
parts.append(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")
return " | ".join(parts), truncated
def _recursive_smart_truncate(obj: Any, max_len: int) -> tuple[Any, bool]:
"""Recursively truncates string values within a dict or list.
Args:
obj: The object to truncate.
max_len: Maximum length for string values.
Returns:
A tuple of (truncated_object, is_truncated).
"""
if isinstance(obj, str):
if max_len != -1 and len(obj) > max_len:
return obj[:max_len] + "...[TRUNCATED]", True
return obj, False
elif isinstance(obj, dict):
truncated_any = False
# Use dict comprehension for potentially slightly better performance,
# but explicit loop is fine for clarity given recursive nature.
new_dict = {}
for k, v in obj.items():
val, trunc = _recursive_smart_truncate(v, max_len)
if trunc:
truncated_any = True
new_dict[k] = val
return new_dict, truncated_any
elif isinstance(obj, (list, tuple)):
truncated_any = False
new_list = []
# Explicit loop to handle flag propagation
for i in obj:
val, trunc = _recursive_smart_truncate(i, max_len)
if trunc:
truncated_any = True
new_list.append(val)
return type(obj)(new_list), truncated_any
elif dataclasses.is_dataclass(obj) and not isinstance(obj, type):
# Convert dataclasses to dicts so they become valid JSON objects
return _recursive_smart_truncate(dataclasses.asdict(obj), max_len)
elif hasattr(obj, "model_dump") and callable(obj.model_dump):
# Pydantic v2
try:
return _recursive_smart_truncate(obj.model_dump(), max_len)
except Exception:
pass
elif hasattr(obj, "dict") and callable(obj.dict):
# Pydantic v1
try:
return _recursive_smart_truncate(obj.dict(), max_len)
except Exception:
pass
elif hasattr(obj, "to_dict") and callable(obj.to_dict):
# Common pattern for custom objects
try:
return _recursive_smart_truncate(obj.to_dict(), max_len)
except Exception:
pass
elif obj is None or isinstance(obj, (int, float, bool)):
# Basic types are safe
return obj, False
# Fallback for unknown types: Convert to string to ensure JSON validity
# We return string representation of the object, which is a valid JSON string value.
return str(obj), False
# --- PyArrow Helper Functions ---
def _pyarrow_datetime() -> pa.DataType:
return pa.timestamp("us", tz=None)
def _pyarrow_numeric() -> pa.DataType:
return pa.decimal128(38, 9)
def _pyarrow_bignumeric() -> pa.DataType:
return pa.decimal256(76, 38)
def _pyarrow_time() -> pa.DataType:
return pa.time64("us")
def _pyarrow_timestamp() -> pa.DataType:
return pa.timestamp("us", tz="UTC")
_BQ_TO_ARROW_SCALARS = MappingProxyType({
"BOOL": pa.bool_,
"BOOLEAN": pa.bool_,
"BYTES": pa.binary,
"DATE": pa.date32,
"DATETIME": _pyarrow_datetime,
"FLOAT": pa.float64,
"FLOAT64": pa.float64,
"GEOGRAPHY": pa.string,
"INT64": pa.int64,
"INTEGER": pa.int64,
"JSON": pa.string,
"NUMERIC": _pyarrow_numeric,
"BIGNUMERIC": _pyarrow_bignumeric,
"STRING": pa.string,
"TIME": _pyarrow_time,
"TIMESTAMP": _pyarrow_timestamp,
})
_BQ_FIELD_TYPE_TO_ARROW_FIELD_METADATA = {
"GEOGRAPHY": {
b"ARROW:extension:name": b"google:sqlType:geography",
b"ARROW:extension:metadata": b'{"encoding": "WKT"}',
},
"DATETIME": {b"ARROW:extension:name": b"google:sqlType:datetime"},
"JSON": {b"ARROW:extension:name": b"google:sqlType:json"},
}
_STRUCT_TYPES = ("RECORD", "STRUCT")
def _bq_to_arrow_scalars(bq_scalar: str) -> Optional[Callable[[], pa.DataType]]:
"""Maps BigQuery scalar types to PyArrow type constructors."""
return _BQ_TO_ARROW_SCALARS.get(bq_scalar)
def _bq_to_arrow_field(bq_field: bq_schema.SchemaField) -> Optional[pa.Field]:
"""Converts a BigQuery SchemaField to a PyArrow Field."""
arrow_type = _bq_to_arrow_data_type(bq_field)
if arrow_type:
metadata = _BQ_FIELD_TYPE_TO_ARROW_FIELD_METADATA.get(
bq_field.field_type.upper() if bq_field.field_type else ""
)
nullable = bq_field.mode.upper() != "REQUIRED"
return pa.field(
bq_field.name, arrow_type, nullable=nullable, metadata=metadata
)
logger.warning(
"Could not determine Arrow type for field '%s' with type '%s'.",
bq_field.name,
bq_field.field_type,
)
return None
def _bq_to_arrow_struct_data_type(
field: bq_schema.SchemaField,
) -> Optional[pa.StructType]:
"""Converts a BigQuery RECORD/STRUCT field to a PyArrow StructType."""
arrow_fields = []
for subfield in field.fields:
arrow_subfield = _bq_to_arrow_field(subfield)
if arrow_subfield:
arrow_fields.append(arrow_subfield)
else:
logger.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_data_type(
field: bq_schema.SchemaField,
) -> Optional[pa.DataType]:
"""Converts a BigQuery field to a PyArrow DataType."""
if field.mode == "REPEATED":
inner = _bq_to_arrow_data_type(
bq_schema.SchemaField(field.name, field.field_type, fields=field.fields)
)
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)
constructor = _bq_to_arrow_scalars(field_type_upper)
if constructor:
return constructor()
else:
logger.warning(
"Failed to convert BigQuery field '%s': unsupported type '%s'.",
field.name,
field.field_type,
)
return None
def to_arrow_schema(
bq_schema_list: list[bq_schema.SchemaField],
) -> Optional[pa.Schema]:
"""Converts a list of BigQuery SchemaFields to a PyArrow Schema.
Args:
bq_schema_list: list of bigquery.SchemaField objects.
Returns:
pa.Schema or None if conversion fails.
"""
arrow_fields = []
for bq_field in bq_schema_list:
af = _bq_to_arrow_field(bq_field)
if af:
arrow_fields.append(af)
else:
logger.error("Failed to convert schema due to field '%s'.", bq_field.name)
return None
return pa.schema(arrow_fields)
# ==============================================================================
# CONFIGURATION
# ==============================================================================
@dataclass
class RetryConfig:
"""Configuration for retrying failed BigQuery write operations.
Attributes:
max_retries: Maximum number of retry attempts.
initial_delay: Initial delay between retries in seconds.
multiplier: Multiplier for exponential backoff.
max_delay: Maximum delay between retries in seconds.
"""
max_retries: int = 3
initial_delay: float = 1.0
multiplier: float = 2.0
max_delay: float = 10.0
@dataclass
class BigQueryLoggerConfig:
"""Configuration for the BigQueryAgentAnalyticsPlugin.
Attributes:
enabled: Whether logging is enabled.
event_allowlist: list of event types to log. If None, all are allowed.
event_denylist: list of event types to ignore.
max_content_length: Max length for text content before truncation.
table_id: BigQuery table ID.
clustering_fields: Fields to cluster the table by.
log_multi_modal_content: Whether to log detailed content parts.
retry_config: Retry configuration for writes.
batch_size: Number of rows per batch.
batch_flush_interval: Max time to wait before flushing a batch.
shutdown_timeout: Max time to wait for shutdown.
queue_max_size: Max size of the in-memory queue.
content_formatter: Optional custom formatter for content.
"""
enabled: bool = True
# V1 Configuration Parity
event_allowlist: list[str] | None = None
event_denylist: list[str] | None = None
max_content_length: int = 500 * 1024 # Defaults to 500KB per text block
table_id: str = "agent_events_v2"
# V2 Configuration
clustering_fields: list[str] = field(
default_factory=lambda: ["event_type", "agent", "user_id"]
)
log_multi_modal_content: bool = True
retry_config: RetryConfig = field(default_factory=RetryConfig)
batch_size: int = 1
batch_flush_interval: float = 1.0
shutdown_timeout: float = 10.0
queue_max_size: int = 10000
content_formatter: Optional[Callable[[Any, str], Any]] = None
# If provided, large content (images, audio, video, large text) will be offloaded to this GCS bucket.
gcs_bucket_name: Optional[str] = None
# If provided, this connection ID will be used as the authorizer for ObjectRef columns.
# Format: "location.connection_id" (e.g. "us.my-connection")
connection_id: Optional[str] = None
# ==============================================================================
# HELPER: TRACE MANAGER (Async-Safe with ContextVars)
# ==============================================================================
_trace_id_ctx = contextvars.ContextVar("_bq_analytics_trace_id", default=None)
_root_agent_name_ctx = contextvars.ContextVar(
"_bq_analytics_root_agent_name", default=None
)
_span_stack_ctx = contextvars.ContextVar("_bq_analytics_span_stack", default=())
_span_times_ctx = contextvars.ContextVar(
"_bq_analytics_span_times", default=None
)
_span_first_token_times_ctx = contextvars.ContextVar(
"_bq_analytics_span_first_token_times", default=None
)
class TraceManager:
"""Manages OpenTelemetry-style trace and span context using contextvars."""
@staticmethod
def init_trace(callback_context: CallbackContext) -> None:
if _trace_id_ctx.get() is None:
_trace_id_ctx.set(callback_context.invocation_id)
# Extract root agent name from invocation context
try:
root_agent = callback_context._invocation_context.agent.root_agent
_root_agent_name_ctx.set(root_agent.name)
except (AttributeError, ValueError):
pass
_span_stack_ctx.set(())
_span_times_ctx.set({})
_span_first_token_times_ctx.set({})
@staticmethod
def get_trace_id(callback_context: CallbackContext) -> Optional[str]:
# Try contextvars first
if trace_id := _trace_id_ctx.get():
return trace_id
# Fallback to callback_context for existing tests/legacy flows
return callback_context.state.get("_bq_analytics_trace_id")
@staticmethod
def push_span(
callback_context: CallbackContext, span_id: Optional[str] = None
) -> str:
# Ensure trace is initialized
if _trace_id_ctx.get() is None:
TraceManager.init_trace(callback_context)
span_id = span_id or str(uuid.uuid4())
stack = _span_stack_ctx.get()
new_stack = stack + (span_id,)
_span_stack_ctx.set(new_stack)
times = dict(_span_times_ctx.get() or {})
times[span_id] = time.time()
_span_times_ctx.set(times)
return span_id
@staticmethod
def pop_span() -> tuple[Optional[str], Optional[int]]:
stack = list(_span_stack_ctx.get())
if not stack:
return None, None
span_id = stack.pop()
_span_stack_ctx.set(tuple(stack))
times_dict = dict(_span_times_ctx.get() or {})
start_time = times_dict.pop(span_id, None)
_span_times_ctx.set(times_dict)
ft_dict = dict(_span_first_token_times_ctx.get() or {})
ft_dict.pop(span_id, None)
_span_first_token_times_ctx.set(ft_dict)
duration_ms = int((time.time() - start_time) * 1000) if start_time else None
return span_id, duration_ms
@staticmethod
def get_current_span_and_parent() -> tuple[Optional[str], Optional[str]]:
stack = _span_stack_ctx.get()
if not stack:
return None, None
return stack[-1], (stack[-2] if len(stack) > 1 else None)
@staticmethod
def get_current_span_id() -> Optional[str]:
stack = _span_stack_ctx.get()
return stack[-1] if stack else None
@staticmethod
def get_root_agent_name() -> Optional[str]:
return _root_agent_name_ctx.get()
@staticmethod
def get_start_time(span_id: str) -> Optional[float]:
times = _span_times_ctx.get()
return times.get(span_id) if times else None
@staticmethod
def record_first_token(span_id: str) -> bool:
"""Records the current time as first token time if not already recorded.
Returns:
True if this was the first token (newly recorded), False otherwise.
"""
first_tokens = dict(_span_first_token_times_ctx.get() or {})
if span_id not in first_tokens:
first_tokens[span_id] = time.time()
_span_first_token_times_ctx.set(first_tokens)
return True
return False
@staticmethod
def get_first_token_time(span_id: str) -> Optional[float]:
first_tokens = _span_first_token_times_ctx.get()
return first_tokens.get(span_id) if first_tokens else None
# ==============================================================================
# HELPER: BATCH PROCESSOR
# ==============================================================================
class BatchProcessor:
"""Handles asynchronous batching and writing of events to BigQuery."""
def __init__(
self,
write_client: BigQueryWriteAsyncClient,
arrow_schema: pa.Schema,
write_stream: str,
batch_size: int,
flush_interval: float,
retry_config: RetryConfig,
queue_max_size: int,
shutdown_timeout: float,
):
"""Initializes the BatchProcessor.
Args:
write_client: BigQueryWriteAsyncClient for writing rows.
arrow_schema: PyArrow schema for serialization.
write_stream: BigQuery write stream name.
batch_size: Number of rows per batch.
flush_interval: Max time to wait before flushing a batch.
retry_config: Retry configuration.
queue_max_size: Max size of the in-memory queue.
shutdown_timeout: Max time to wait for shutdown.
"""
self.write_client = write_client
self.arrow_schema = arrow_schema
self.write_stream = write_stream
self.batch_size = batch_size
self.flush_interval = flush_interval
self.retry_config = retry_config
self.shutdown_timeout = shutdown_timeout
self._queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(
maxsize=queue_max_size
)
self._batch_processor_task: Optional[asyncio.Task] = None
self._shutdown = False
async def start(self):
"""Starts the batch writer worker task."""
if self._batch_processor_task is None:
self._batch_processor_task = asyncio.create_task(self._batch_writer())
async def append(self, row: dict[str, Any]) -> None:
"""Appends a row to the queue for batching.
Args:
row: Dictionary representing a single row.
"""
try:
self._queue.put_nowait(row)
except asyncio.QueueFull:
logger.warning("BigQuery log queue full, dropping event.")
def _prepare_arrow_batch(self, rows: list[dict[str, Any]]) -> pa.RecordBatch:
"""Prepares a PyArrow RecordBatch from a list of rows.
Args:
rows: list of row dictionaries.
Returns:
pa.RecordBatch for writing.
"""
data = {field.name: [] for field in self.arrow_schema}
for row in rows:
for field in self.arrow_schema:
value = row.get(field.name)
# JSON fields must be serialized to strings for the Arrow layer
field_metadata = self.arrow_schema.field(field.name).metadata
is_json = False
if field_metadata and b"ARROW:extension:name" in field_metadata:
if field_metadata[b"ARROW:extension:name"] == b"google:sqlType:json":
is_json = True
arrow_field_type = self.arrow_schema.field(field.name).type
is_struct = pa.types.is_struct(arrow_field_type)
is_list = pa.types.is_list(arrow_field_type)
if is_json:
if value is not None:
if isinstance(value, (dict, list)):
try:
value = json.dumps(value)
except (TypeError, ValueError):
value = str(value)
elif isinstance(value, (str, bytes)):
if isinstance(value, bytes):
try:
value = value.decode("utf-8")
except UnicodeDecodeError:
value = str(value)
# Check if it's already a valid JSON object or array to avoid double-encoding
is_already_json = False
if isinstance(value, str):
stripped = value.strip()
if stripped.startswith(("{", "[")) and stripped.endswith(
("}", "]")
):
try:
json.loads(value)
is_already_json = True
except (ValueError, TypeError):
pass
if not is_already_json:
try:
value = json.dumps(value)
except (TypeError, ValueError):
value = str(value)
# If is_already_json is True, we keep value as-is
else:
# For other types (int, float, bool), serialize to JSON equivalents
try:
value = json.dumps(value)
except (TypeError, ValueError):
value = str(value)
elif isinstance(value, (dict, list)) and not is_struct and not is_list:
if value is not None and not isinstance(value, (str, bytes)):
try:
value = json.dumps(value)
except (TypeError, ValueError):
value = str(value)
data[field.name].append(value)
return pa.RecordBatch.from_pydict(data, schema=self.arrow_schema)
async def _batch_writer(self) -> None:
"""Worker task that batches and writes rows to BigQuery."""
while not self._shutdown or not self._queue.empty():
batch = []
try:
if self._shutdown:
try:
first_item = self._queue.get_nowait()
except asyncio.QueueEmpty:
break
else:
first_item = await asyncio.wait_for(
self._queue.get(), timeout=self.flush_interval
)
batch.append(first_item)
while len(batch) < self.batch_size:
try:
item = self._queue.get_nowait()
batch.append(item)
except asyncio.QueueEmpty:
break
if batch:
try:
await self._write_rows_with_retry(batch)
finally:
# Mark tasks as done ONLY after processing (write attempt)
for _ in batch:
self._queue.task_done()
except asyncio.TimeoutError:
continue
except asyncio.CancelledError:
logger.info("Batch writer task cancelled.")
break
except Exception as e:
logger.error("Error in batch writer loop: %s", e, exc_info=True)
await asyncio.sleep(1)
async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None:
"""Writes a batch of rows to BigQuery with retry logic.
Args:
rows: list of row dictionaries to write.
"""
attempt = 0
delay = self.retry_config.initial_delay
try:
arrow_batch = self._prepare_arrow_batch(rows)
serialized_schema = self.arrow_schema.serialize().to_pybytes()
serialized_batch = arrow_batch.serialize().to_pybytes()
req = bq_storage_types.AppendRowsRequest(
write_stream=self.write_stream,
trace_id=f"google-adk-bq-logger/{__version__}",
)
req.arrow_rows.writer_schema.serialized_schema = serialized_schema
req.arrow_rows.rows.serialized_record_batch = serialized_batch
except Exception as e:
logger.error(
"Failed to prepare Arrow batch (Data Loss): %s", e, exc_info=True
)
return
while attempt <= self.retry_config.max_retries:
try:
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:
if "schema mismatch" in error_message.lower():
logger.error(
"BigQuery Schema Mismatch: %s. This usually means the"
" table schema does not match the expected schema.",
error_message,
)
else:
logger.error("Non-retryable BigQuery error: %s", error_message)
row_errors = getattr(response, "row_errors", [])
if row_errors:
for row_error in row_errors:
logger.error("Row error details: %s", row_error)
logger.error("Row content causing error: %s", rows)
return
return
except (ServiceUnavailable, TooManyRequests, InternalServerError) as e:
attempt += 1
if attempt > self.retry_config.max_retries:
logger.error(
"BigQuery Batch Dropped after %s attempts. Last error: %s",
self.retry_config.max_retries + 1,
e,
)
return
sleep_time = min(
delay * (1 + random.random()), self.retry_config.max_delay
)
logger.warning(
"BigQuery write failed (Attempt %s), retrying in %.2fs..."
" Error: %s",
attempt,
sleep_time,
e,
)
await asyncio.sleep(sleep_time)
delay *= self.retry_config.multiplier
except Exception as e:
logger.error(
"Unexpected BigQuery Write API error (Dropping batch): %s",
e,
exc_info=True,
)
return
async def shutdown(self, timeout: float = 5.0) -> None:
"""Shuts down the BatchProcessor, draining the queue.
Args:
timeout: Maximum time to wait for the queue to drain.
"""
self._shutdown = True
logger.info("BatchProcessor shutting down, draining queue...")
if self._batch_processor_task:
try:
await asyncio.wait_for(self._batch_processor_task, timeout=timeout)
except asyncio.TimeoutError:
logger.warning("BatchProcessor shutdown timed out, cancelling worker.")
self._batch_processor_task.cancel()
except Exception as e:
logger.error("Error during BatchProcessor shutdown: %s", e)
async def close(self) -> None:
"""Closes the processor and flushes remaining items."""
if self._shutdown:
return
self._shutdown = True
# Wait for queue to be empty
try:
await asyncio.wait_for(self._queue.join(), timeout=self.shutdown_timeout)
except (asyncio.TimeoutError, asyncio.CancelledError):
logger.warning(
"Timeout waiting for BigQuery batch queue to empty on shutdown."
)
# Cancel the writer task if it's still running (it should exit on _shutdown + empty queue)
if self._batch_processor_task and not self._batch_processor_task.done():
self._batch_processor_task.cancel()
try:
await self._batch_processor_task
except asyncio.CancelledError:
pass
# ==============================================================================
# HELPER: CONTENT PARSER (Length Limits Only)
# ==============================================================================
class ContentParser:
"""Parses content for logging with length limits and structure normalization."""
def __init__(self, max_length: int) -> None:
"""Initializes the ContentParser.
Args:
max_length: Maximum length for text content.
"""
self.max_length = max_length
def _truncate(self, text: str) -> tuple[str, bool]:
if self.max_length != -1 and text and len(text) > self.max_length:
return text[: self.max_length] + "...[TRUNCATED]", True
return text, False
class GCSOffloader:
"""Offloads content to GCS."""
def __init__(
self,
project_id: str,
bucket_name: str,
executor: ThreadPoolExecutor,
storage_client: Optional[storage.Client] = None,
):
self.client = storage_client or storage.Client(project=project_id)
self.bucket = self.client.bucket(bucket_name)
self.executor = executor
async def upload_content(
self, data: bytes | str, content_type: str, path: str
) -> str:
"""Async wrapper around blocking GCS upload."""
loop = asyncio.get_running_loop()
return await loop.run_in_executor(
self.executor,
functools.partial(self._upload_sync, data, content_type, path),
)
def _upload_sync(
self, data: bytes | str, content_type: str, path: str
) -> str:
blob = self.bucket.blob(path)
blob.upload_from_string(data, content_type=content_type)
return f"gs://{self.bucket.name}/{path}"
class HybridContentParser:
"""Parses content and offloads large/binary parts to GCS."""
def __init__(
self,
offloader: Optional[GCSOffloader],
trace_id: str,
span_id: str,
max_length: int = 20000,
connection_id: Optional[str] = None,
):
self.offloader = offloader
self.trace_id = trace_id
self.span_id = span_id
self.max_length = max_length
self.connection_id = connection_id
self.inline_text_limit = 32 * 1024 # 32KB limit
def _truncate(self, text: str) -> tuple[str, bool]:
if self.max_length != -1 and len(text) > self.max_length:
return (
text[: self.max_length] + "...[TRUNCATED]",
True,
)
return text, False
async def _parse_content_object(
self, content: types.Content | types.Part
) -> tuple[str, list[dict[str, Any]], bool]:
"""Parses a Content or Part object into summary text and content parts."""
content_parts = []
is_truncated = False
summary_text = []
parts = content.parts if hasattr(content, "parts") else [content]
for idx, part in enumerate(parts):
part_data = {
"part_index": idx,
"mime_type": "text/plain",
"uri": None,
"text": None,
"part_attributes": "{}",
"storage_mode": "INLINE",
"object_ref": None,
}
# CASE A: It is already a URI (e.g. from user input)
if hasattr(part, "file_data") and part.file_data:
part_data["storage_mode"] = "EXTERNAL_URI"
part_data["uri"] = part.file_data.file_uri
part_data["mime_type"] = part.file_data.mime_type
# CASE B: It is Binary/Inline Data (Image/Blob)
elif hasattr(part, "inline_data") and part.inline_data:
if self.offloader:
ext = mimetypes.guess_extension(part.inline_data.mime_type) or ".bin"
path = f"{datetime.now().date()}/{self.trace_id}/{self.span_id}_p{idx}{ext}"
try:
uri = await self.offloader.upload_content(
part.inline_data.data, part.inline_data.mime_type, path
)
part_data["storage_mode"] = "GCS_REFERENCE"
part_data["uri"] = uri
object_ref = {
"uri": uri,
"version": None,
"authorizer": self.connection_id,
"details": json.dumps({
"gcs_metadata": {"content_type": part.inline_data.mime_type}
}),
}
part_data["object_ref"] = object_ref
part_data["mime_type"] = part.inline_data.mime_type
part_data["text"] = "[MEDIA OFFLOADED]"
except Exception as e:
logger.warning("Failed to offload content to GCS: %s", e)
part_data["text"] = "[UPLOAD FAILED]"
else:
part_data["text"] = "[BINARY DATA]"
# CASE C: Text
elif hasattr(part, "text") and part.text:
text_len = len(part.text.encode("utf-8"))
# If max_length is set and smaller than inline limit, use it as threshold
# to prefer offloading over truncation.
offload_threshold = self.inline_text_limit
if self.max_length != -1 and self.max_length < offload_threshold:
offload_threshold = self.max_length
if self.offloader and text_len > offload_threshold:
# Text is too big, treat as file
path = f"{datetime.now().date()}/{self.trace_id}/{self.span_id}_p{idx}.txt"
try:
uri = await self.offloader.upload_content(
part.text, "text/plain", path
)
part_data["storage_mode"] = "GCS_REFERENCE"
part_data["uri"] = uri
object_ref = {
"uri": uri,
"version": None,
"authorizer": self.connection_id,
"details": json.dumps(
{"gcs_metadata": {"content_type": "text/plain"}}
),
}
part_data["object_ref"] = object_ref
part_data["mime_type"] = "text/plain"
part_data["text"] = part.text[:200] + "... [OFFLOADED]"
except Exception as e:
logger.warning("Failed to offload text to GCS: %s", e)
clean_text, truncated = self._truncate(part.text)
if truncated:
is_truncated = True
part_data["text"] = clean_text
summary_text.append(clean_text)
else:
# Text is small or no offloader, keep inline
clean_text, truncated = self._truncate(part.text)
if truncated:
is_truncated = True
part_data["text"] = clean_text
summary_text.append(clean_text)
elif hasattr(part, "function_call") and part.function_call:
part_data["mime_type"] = "application/json"
part_data["text"] = f"Function: {part.function_call.name}"
part_data["part_attributes"] = json.dumps(
{"function_name": part.function_call.name}
)
content_parts.append(part_data)
summary_str, truncated = self._truncate(" | ".join(summary_text))
if truncated:
is_truncated = True
return summary_str, content_parts, is_truncated
async def parse(self, content: Any) -> tuple[Any, list[dict[str, Any]], bool]:
"""Parses content into JSON payload and content parts, potentially offloading to GCS."""
json_payload = {}
content_parts = []
is_truncated = False
def process_text(t: str) -> tuple[str, bool]:
return self._truncate(t)
if isinstance(content, LlmRequest):
# Handle Prompt
messages = []
contents = (
content.contents
if isinstance(content.contents, list)
else [content.contents]
)
for c in contents:
role = getattr(c, "role", "unknown")
summary, parts, trunc = await self._parse_content_object(c)
if trunc:
is_truncated = True
content_parts.extend(parts)
messages.append({"role": role, "content": summary})
if messages:
json_payload["prompt"] = messages
# Handle System Instruction
if content.config and getattr(content.config, "system_instruction", None):
si = content.config.system_instruction
if isinstance(si, str):
json_payload["system_prompt"] = si
else:
summary, parts, trunc = await self._parse_content_object(si)
if trunc:
is_truncated = True
content_parts.extend(parts)
json_payload["system_prompt"] = summary
elif isinstance(content, (types.Content, types.Part)):
summary, parts, trunc = await self._parse_content_object(content)
return {"text_summary": summary}, parts, trunc
elif isinstance(content, (dict, list)):
json_payload, is_truncated = _recursive_smart_truncate(
content, self.max_length
)
elif isinstance(content, str):
json_payload, is_truncated = process_text(content)
elif content is None:
json_payload = None
else:
json_payload, is_truncated = process_text(str(content))
return json_payload, content_parts, is_truncated
def _get_events_schema() -> list[bigquery.SchemaField]:
"""Returns the BigQuery schema for the events table."""
return [
bigquery.SchemaField(
"timestamp",
"TIMESTAMP",
mode="REQUIRED",
description=(
"The UTC timestamp when the event occurred. Used for ordering"
" events within a session."
),
),
bigquery.SchemaField(
"event_type",
"STRING",
mode="NULLABLE",
description=(
"The category of the event (e.g., 'LLM_REQUEST', 'TOOL_CALL',"
" 'AGENT_RESPONSE'). Helps in filtering specific types of"
" interactions."
),
),
bigquery.SchemaField(
"agent",
"STRING",
mode="NULLABLE",
description=(
"The name of the agent that generated this event. Useful for"
" multi-agent systems."
),
),
bigquery.SchemaField(
"session_id",
"STRING",
mode="NULLABLE",
description=(
"A unique identifier for the entire conversation session. Used"
" to group all events belonging to a single user interaction."
),
),
bigquery.SchemaField(
"invocation_id",
"STRING",
mode="NULLABLE",
description=(
"A unique identifier for a single turn or execution within a"
" session. Groups related events like LLM request and response."
),
),
bigquery.SchemaField(
"user_id",
"STRING",
mode="NULLABLE",
description=(
"The identifier of the end-user participating in the session,"
" if available."
),
),
bigquery.SchemaField(
"trace_id",
"STRING",
mode="NULLABLE",
description=(
"OpenTelemetry trace ID for distributed tracing across services."
),
),
bigquery.SchemaField(
"span_id",
"STRING",
mode="NULLABLE",
description="OpenTelemetry span ID for this specific operation.",
),
bigquery.SchemaField(
"parent_span_id",
"STRING",
mode="NULLABLE",
description=(
"OpenTelemetry parent span ID to reconstruct the operation"
" hierarchy."
),
),
bigquery.SchemaField(
"content",
"JSON",
mode="NULLABLE",
description=(
"The primary payload of the event, stored as a JSON string. The"
" structure depends on the event_type (e.g., prompt text for"
" LLM_REQUEST, tool output for TOOL_RESPONSE)."
),
),
bigquery.SchemaField(
"content_parts",
"RECORD",
mode="REPEATED",
fields=[
bigquery.SchemaField(
"mime_type",
"STRING",
mode="NULLABLE",
description=(
"The MIME type of the content part (e.g., 'text/plain',"
" 'image/png')."
),
),
bigquery.SchemaField(
"uri",
"STRING",
mode="NULLABLE",
description=(
"The URI of the content part if stored externally"
" (e.g., GCS bucket path)."
),
),
bigquery.SchemaField(
"object_ref",
"RECORD",
mode="NULLABLE",
fields=[
bigquery.SchemaField(
"uri",
"STRING",
mode="NULLABLE",
description="The URI of the object.",
),
bigquery.SchemaField(
"version",
"STRING",
mode="NULLABLE",
description="The version of the object.",
),
bigquery.SchemaField(
"authorizer",
"STRING",
mode="NULLABLE",
description="The authorizer for the object.",
),
bigquery.SchemaField(
"details",
"JSON",
mode="NULLABLE",
description="Additional details about the object.",
),
],
description=(
"The ObjectRef of the content part if stored externally."
),
),
bigquery.SchemaField(
"text",
"STRING",
mode="NULLABLE",
description="The raw text content if the part is text-based.",
),
bigquery.SchemaField(
"part_index",
"INTEGER",
mode="NULLABLE",
description=(
"The zero-based index of this part within the content."
),
),
bigquery.SchemaField(
"part_attributes",
"STRING",
mode="NULLABLE",
description=(
"Additional metadata for this content part as a JSON"
" object (serialized to string)."
),
),
bigquery.SchemaField(
"storage_mode",
"STRING",
mode="NULLABLE",
description=(
"Indicates how the content part is stored (e.g.,"
" 'INLINE', 'GCS_REFERENCE', 'EXTERNAL_URI')."
),
),
],
description=(
"For multi-modal events, contains a list of content parts"
" (text, images, etc.)."
),
),
bigquery.SchemaField(
"attributes",
"JSON",
mode="NULLABLE",
description=(
"A JSON object containing arbitrary key-value pairs for"
" additional event metadata. Includes enrichment fields like"
" 'root_agent_name' (turn orchestration), 'model' (request"
" model), 'model_version' (response version), and"
" 'usage_metadata' (detailed token counts)."
),
),
bigquery.SchemaField(
"latency_ms",
"JSON",
mode="NULLABLE",
description=(
"A JSON object containing latency measurements, such as"
" 'total_ms' and 'time_to_first_token_ms'."
),
),
bigquery.SchemaField(
"status",
"STRING",
mode="NULLABLE",
description="The outcome of the event, typically 'OK' or 'ERROR'.",
),
bigquery.SchemaField(
"error_message",
"STRING",
mode="NULLABLE",
description="Detailed error message if the status is 'ERROR'.",
),
bigquery.SchemaField(
"is_truncated",
"BOOLEAN",
mode="NULLABLE",
description=(
"Boolean flag indicating if the 'content' field was truncated"
" because it exceeded the maximum allowed size."
),
),
]
# ==============================================================================
# MAIN PLUGIN
# ==============================================================================
_GLOBAL_WRITE_CLIENT: Optional[BigQueryWriteAsyncClient] = None
_GLOBAL_CLIENT_LOCK = asyncio.Lock()
class BigQueryAgentAnalyticsPlugin(BasePlugin):
"""BigQuery Agent Analytics Plugin (v2.0 using Write API).
Logs agent events (LLM requests, tool calls, etc.) to BigQuery for analytics.
Uses the BigQuery Write API for efficient, asynchronous, and reliable logging.
"""
def __init__(
self,
project_id: str,
dataset_id: str,
*,
table_id: Optional[str] = None,
config: Optional[BigQueryLoggerConfig] = None,
location: str = "US",
) -> None:
"""Initializes the BigQueryAgentAnalyticsPlugin.
Args:
project_id: Google Cloud project ID.
dataset_id: BigQuery dataset ID.
table_id: BigQuery table ID (optional, overrides config).
config: BigQueryLoggerConfig (optional).
location: BigQuery location (default: "US").
"""
super().__init__(name="bigquery_agent_analytics")
self.project_id = project_id
self.dataset_id = dataset_id
self.config = config or BigQueryLoggerConfig()
self.table_id = table_id or self.config.table_id
self.location = location
self._started = False
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._executor = None
self.offloader: Optional[GCSOffloader] = None
self.parser: Optional[HybridContentParser] = None
def _format_content_safely(
self, content: Optional[types.Content]
) -> tuple[str, bool]:
"""Formats content using config.content_formatter or default formatter.
Args:
content: The content to format.
Returns:
A tuple of (formatted_string, is_truncated).
"""
if content is None:
return "None", False
try:
# If a custom formatter is provided, we could try to use it here too,
# but it expects (content, event_type). For internal formatting,
# we stick to the default _format_content but respect max_len.
return _format_content(content, max_len=self.config.max_content_length)
except Exception as e:
logger.warning("Content formatter failed: %s", e)
return "[FORMATTING FAILED]", False
async def _lazy_setup(self, **kwargs) -> None:
"""Performs lazy initialization of BigQuery clients and resources."""
if self._started:
return
loop = asyncio.get_running_loop()
if not self.client:
if self._executor is None:
self._executor = ThreadPoolExecutor(max_workers=1)
self.client = await loop.run_in_executor(
self._executor,
lambda: bigquery.Client(
project=self.project_id, location=self.location
),
)
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.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:
self.arrow_schema = to_arrow_schema(self._schema)
if not self.arrow_schema:
raise RuntimeError("Failed to convert BigQuery schema to Arrow schema.")
self.offloader = None
if self.config.gcs_bucket_name:
self.offloader = GCSOffloader(
self.project_id,
self.config.gcs_bucket_name,
self._executor,
storage_client=kwargs.get("storage_client"),
)
self.parser = HybridContentParser(
self.offloader,
"",
"",
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))
@staticmethod
def _atexit_cleanup(batch_processor: "BatchProcessor") -> None:
"""Clean up batch processor on script exit."""
# Check if the batch_processor object is still alive
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)
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:
self.client.get_table(self.full_table_id)
except cloud_exceptions.NotFound:
logger.info("Table %s not found, creating table.", self.full_table_id)
tbl = bigquery.Table(self.full_table_id, schema=self._schema)
tbl.time_partitioning = bigquery.TimePartitioning(
type_=bigquery.TimePartitioningType.DAY, field="timestamp"
)
tbl.clustering_fields = self.config.clustering_fields
try:
self.client.create_table(tbl)
except cloud_exceptions.Conflict:
pass
except Exception as e:
logger.error(
"Could not create table %s: %s",
self.full_table_id,
e,
exc_info=True,
)
except Exception as e:
logger.error(
"Error checking for table %s: %s",
self.full_table_id,
e,
exc_info=True,
)
async def shutdown(self, timeout: float | None = None) -> None:
"""Shuts down the plugin and releases resources.
Args:
timeout: Maximum time to wait for the queue to drain.
"""
if self._is_shutting_down:
return
self._is_shutting_down = True
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()
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
self._started = False
def __getstate__(self):
"""Custom pickling to exclude non-picklable runtime objects."""
state = self.__dict__.copy()
state["_setup_lock"] = None
state["client"] = None
state["write_client"] = None
state["write_stream"] = None
state["batch_processor"] = None
state["_executor"] = None
state["offloader"] = None
state["parser"] = None
state["_started"] = False
state["_is_shutting_down"] = False
return state
def __setstate__(self, state):
"""Custom unpickling to restore state."""
self.__dict__.update(state)
async def __aenter__(self) -> BigQueryAgentAnalyticsPlugin:
await self._ensure_started()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
await self.shutdown()
async def _ensure_started(self, **kwargs) -> None:
"""Ensures that the plugin is started and initialized."""
if not self._started:
# Kept original lock name as it was not explicitly changed.
if self._setup_lock is None:
self._setup_lock = asyncio.Lock()
async with self._setup_lock:
if not self._started:
try:
await self._lazy_setup(**kwargs)
self._started = True
except Exception as e:
logger.error("Failed to initialize BigQuery Plugin: %s", e)
async def _log_event(
self,
event_type: str,
callback_context: CallbackContext,
raw_content: Any = None,
is_truncated: bool = False,
**kwargs,
) -> None:
"""Logs an event to BigQuery.
Args:
event_type: The type of event (e.g., 'LLM_REQUEST').
callback_context: The callback context.
raw_content: The raw content to log.
is_truncated: Whether the content is already truncated.
**kwargs: Additional attributes to log.
"""
if not self.config.enabled or self._is_shutting_down:
return
if self.config.event_denylist and event_type in self.config.event_denylist:
return
if (
self.config.event_allowlist
and event_type not in self.config.event_allowlist
):
return
if not self._started:
await self._ensure_started()
if not self._started:
return
timestamp = datetime.now(timezone.utc)
if self.config.content_formatter:
try:
raw_content = self.config.content_formatter(raw_content, event_type)
except Exception as e:
logger.warning("Content formatter failed: %s", e)
trace_id = TraceManager.get_trace_id(callback_context)
current_span_id, current_parent_span_id = (
TraceManager.get_current_span_and_parent()
)
span_id = current_span_id
if "span_id_override" in kwargs:
val = kwargs.pop("span_id_override")
if val is not None:
span_id = val
parent_span_id = current_parent_span_id
if "parent_span_id_override" in kwargs:
val = kwargs.pop("parent_span_id_override")
if val is not None:
parent_span_id = val
# Use HybridContentParser if offloader is available, otherwise use default
# Re-initialize parser with current trace/span IDs for GCS pathing
self.parser = HybridContentParser(
self.offloader,
trace_id or "no_trace",
span_id or "no_span",
self.config.max_content_length,
connection_id=self.config.connection_id,
)
content_json, content_parts, parser_truncated = await self.parser.parse(
raw_content
)
is_truncated = is_truncated or parser_truncated
total_latency = kwargs.get("latency_ms")
tfft = kwargs.get("time_to_first_token_ms")
latency_json = {}
if total_latency is not None:
latency_json["total_ms"] = total_latency
if tfft is not None:
latency_json["time_to_first_token_ms"] = tfft
kwargs.pop("latency_ms", None)
kwargs.pop("time_to_first_token_ms", None)
# Check if content was truncated by the parser or explicitly passed
# (Already handled by parser_truncated above, but keeping for safety or if other logic added later)
status = kwargs.pop("status", "OK")
error_message = kwargs.pop("error_message", None)
# V2 Metadata Extensions
model = kwargs.pop("model", None)
model_version = kwargs.pop("model_version", None)
usage_metadata = kwargs.pop("usage_metadata", None)
# Add new fields to attributes instead of columns
kwargs["root_agent_name"] = TraceManager.get_root_agent_name()
if model:
kwargs["model"] = model
if model_version:
kwargs["model_version"] = model_version
if usage_metadata:
# Use smart truncate to handle Pydantic, Dataclasses, and other objects
usage_dict, _ = _recursive_smart_truncate(
usage_metadata, self.config.max_content_length
)
if isinstance(usage_dict, dict):
kwargs["usage_metadata"] = usage_dict
else:
# Fallback if it couldn't be converted to dict
kwargs["usage_metadata"] = usage_metadata
# Serialize remaining kwargs to JSON string for attributes
try:
attributes_json = json.dumps(kwargs)
except (TypeError, ValueError):
# Fallback for non-serializable objects
attributes_json = json.dumps(kwargs, default=str)
row = {
"timestamp": timestamp,
"event_type": event_type,
"agent": callback_context.agent_name,
"user_id": callback_context.user_id,
"session_id": callback_context.session.id,
"invocation_id": callback_context.invocation_id,
"trace_id": trace_id,
"span_id": span_id,
"parent_span_id": parent_span_id,
"content": content_json,
"content_parts": (
content_parts if self.config.log_multi_modal_content else []
),
"attributes": attributes_json,
"latency_ms": latency_json if latency_json else None,
"status": status,
"error_message": error_message,
"is_truncated": is_truncated,
}
if self.batch_processor:
await self.batch_processor.append(row)
# --- UPDATED CALLBACKS FOR V1 PARITY ---
async def on_user_message_callback(
self,
*,
invocation_context: InvocationContext,
user_message: types.Content,
**kwargs,
) -> None:
"""Parity with V1: Logs USER_MESSAGE_RECEIVED event.
Args:
invocation_context: The context of the current invocation.
user_message: The message content received from the user.
"""
await self._log_event(
"USER_MESSAGE_RECEIVED",
CallbackContext(invocation_context),
raw_content=user_message,
)
async def before_run_callback(
self, *, invocation_context: "InvocationContext", **kwargs
) -> None:
"""Callback before the agent run starts.
Args:
invocation_context: The context of the current invocation.
"""
await self._ensure_started()
await self._log_event(
"INVOCATION_STARTING", CallbackContext(invocation_context)
)
async def after_run_callback(
self, *, invocation_context: "InvocationContext", **kwargs
) -> None:
"""Callback after the agent run completes.
Args:
invocation_context: The context of the current invocation.
"""
await self._log_event(
"INVOCATION_COMPLETED", CallbackContext(invocation_context)
)
async def before_agent_callback(
self, *, agent: Any, callback_context: CallbackContext, **kwargs
) -> None:
"""Callback before an agent starts processing.
Args:
agent: The agent instance.
callback_context: The callback context.
"""
TraceManager.init_trace(callback_context)
TraceManager.push_span(callback_context)
await self._log_event(
"AGENT_STARTING",
callback_context,
raw_content=getattr(agent, "instruction", ""),
)
async def after_agent_callback(
self, *, agent: Any, callback_context: CallbackContext, **kwargs
) -> None:
"""Callback after an agent completes processing.
Args:
agent: The agent instance.
callback_context: The callback context.
"""
span_id, duration = TraceManager.pop_span()
# When popping, the current stack now points to parent.
# The event we are logging ("AGENT_COMPLETED") belongs to the span we just popped.
# So we must override span_id to be the popped span, and parent to be current top of stack.
parent_span_id, _ = TraceManager.get_current_span_and_parent()
await self._log_event(
"AGENT_COMPLETED",
callback_context,
latency_ms=duration,
span_id_override=span_id,
parent_span_id_override=parent_span_id,
)
async def before_model_callback(
self,
*,
callback_context: CallbackContext,
llm_request: LlmRequest,
**kwargs,
) -> None:
"""Callback before LLM call.
Logs the LLM request details including:
1. Prompt content
2. System instruction (if available)
The content is formatted as 'Prompt: {prompt} | System Prompt:
{system_prompt}'.
"""
# 5. Attributes (Config & Tools)
attributes = {}
if llm_request.config:
config_dict = {}
for field_name in [
"temperature",
"top_p",
"top_k",
"candidate_count",
"max_output_tokens",
"stop_sequences",
]:
if val := getattr(llm_request.config, field_name, None):
config_dict[field_name] = val
if config_dict:
attributes["llm_config"] = config_dict
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)
await self._log_event(
"LLM_REQUEST",
callback_context,
raw_content=llm_request,
model=llm_request.model,
**attributes,
)
async def after_model_callback(
self,
*,
callback_context: CallbackContext,
llm_response: "LlmResponse",
**kwargs,
) -> None:
"""Callback after LLM call.
Logs the LLM response details including:
1. Response content
2. Token usage (if available)
The content is formatted as 'Response: {content} | Usage: {usage}'.
Args:
callback_context: The callback context.
llm_response: The LLM response object.
"""
content_dict = {}
is_truncated = False
if llm_response.content:
part_str, part_truncated = self._format_content_safely(
llm_response.content
)
if part_str:
content_dict["response"] = part_str
if part_truncated:
is_truncated = True
if llm_response.usage_metadata:
usage = llm_response.usage_metadata
usage_dict = {}
if hasattr(usage, "prompt_token_count"):
usage_dict["prompt"] = usage.prompt_token_count
if hasattr(usage, "candidates_token_count"):
usage_dict["completion"] = usage.candidates_token_count
if hasattr(usage, "total_token_count"):
usage_dict["total"] = usage.total_token_count
if usage_dict:
content_dict["usage"] = usage_dict
if content_dict:
content_str = content_dict
else:
content_str = None
span_id = TraceManager.get_current_span_id()
_, parent_span_id = TraceManager.get_current_span_and_parent()
is_popped = False
duration = 0
tfft = None
if hasattr(llm_response, "partial") and llm_response.partial:
# Streaming chunk - do NOT pop span yet
if span_id:
TraceManager.record_first_token(span_id)
start_time = TraceManager.get_start_time(span_id)
first_token = TraceManager.get_first_token_time(span_id)
if start_time:
duration = int((time.time() - start_time) * 1000)
if start_time and first_token:
tfft = int((first_token - start_time) * 1000)
else:
# Final response - pop span
start_time = None
if span_id:
# Ensure we have first token time even if it wasn't streaming (or single chunk)
TraceManager.record_first_token(span_id)
start_time = TraceManager.get_start_time(span_id)
first_token = TraceManager.get_first_token_time(span_id)
if start_time and first_token:
tfft = int((first_token - start_time) * 1000)
# ACTUALLY pop the span
popped_span_id, duration = TraceManager.pop_span()
is_popped = True
# If we popped, the span_id from get_current_span_and_parent() above is correct for THIS event
# Wait, if we popped, get_current_span_and_parent() now returns parent.
# But we captured span_id BEFORE popping. So we should use THAT.
# If is_popped is True, we must override span_id in log_event to use the popped one.
# Otherwise log_event will fetch current stack (which is parent).
span_id = popped_span_id or span_id
extra_kwargs = {}
if tfft is not None:
extra_kwargs["time_to_first_token_ms"] = tfft
await self._log_event(
"LLM_RESPONSE",
callback_context,
raw_content=content_str,
is_truncated=is_truncated,
latency_ms=duration,
model_version=llm_response.model_version,
usage_metadata=llm_response.usage_metadata,
span_id_override=span_id if is_popped else None,
parent_span_id_override=parent_span_id
if is_popped
else None, # Use pre-pop state
**extra_kwargs,
**kwargs,
)
async def on_model_error_callback(
self, *, callback_context: CallbackContext, error: Exception, **kwargs
) -> None:
"""Callback on LLM error.
Args:
callback_context: The callback context.
error: The exception that occurred.
**kwargs: Additional arguments.
"""
span_id, duration = TraceManager.pop_span()
parent_span_id, _ = TraceManager.get_current_span_and_parent()
await self._log_event(
"LLM_ERROR",
callback_context,
error_message=str(error),
latency_ms=duration,
span_id_override=span_id,
parent_span_id_override=parent_span_id,
)
async def before_tool_callback(
self,
*,
tool: BaseTool,
tool_args: dict[str, Any],
tool_context: ToolContext,
**kwargs,
) -> None:
"""Callback before tool execution.
Args:
tool: The tool being executed.
tool_args: The arguments passed to the tool.
tool_context: The tool context.
"""
args_truncated, is_truncated = _recursive_smart_truncate(
tool_args, self.config.max_content_length
)
content_dict = {"tool": tool.name, "args": args_truncated}
TraceManager.push_span(tool_context)
await self._log_event(
"TOOL_STARTING",
tool_context,
raw_content=content_dict,
is_truncated=is_truncated,
)
async def after_tool_callback(
self,
*,
tool: BaseTool,
tool_args: dict[str, Any],
tool_context: ToolContext,
result: dict[str, Any],
**kwargs,
) -> None:
"""Callback after tool execution.
Args:
tool: The tool that was executed.
tool_args: The arguments passed to the tool.
tool_context: The tool context.
result: The response from the tool.
"""
resp_truncated, is_truncated = _recursive_smart_truncate(
result, self.config.max_content_length
)
content_dict = {"tool": tool.name, "result": resp_truncated}
span_id, duration = TraceManager.pop_span()
parent_span_id, _ = TraceManager.get_current_span_and_parent()
await self._log_event(
"TOOL_COMPLETED",
tool_context,
raw_content=content_dict,
is_truncated=is_truncated,
latency_ms=duration,
span_id_override=span_id,
parent_span_id_override=parent_span_id,
)
async def on_tool_error_callback(
self,
*,
tool: BaseTool,
tool_args: dict[str, Any],
tool_context: ToolContext,
error: Exception,
**kwargs,
) -> None:
"""Callback on tool error.
Args:
tool: The tool that failed.
tool_args: The arguments passed to the tool.
tool_context: The tool context.
error: The exception that occurred.
**kwargs: Additional arguments.
"""
args_truncated, is_truncated = _recursive_smart_truncate(
tool_args, self.config.max_content_length
)
content_dict = {"tool": tool.name, "args": args_truncated}
_, duration = TraceManager.pop_span()
await self._log_event(
"TOOL_ERROR",
tool_context,
raw_content=content_dict,
error_message=str(error),
is_truncated=is_truncated,
latency_ms=duration,
)
Reference in New Issue View Git Blame Copy Permalink