pip install "langchain-google-community[bigquery]" langchain langchain-google-genai langgraph

import asyncio
import os
from datetime import datetime

from langchain.tools import tool
from langchain_google_community.callbacks.bigquery_callback import AsyncBigQueryCallbackHandler
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.prebuilt import create_react_agent

# 1. Define some tools for the agent
@tool
def get_current_time():
    """Returns the current local time."""
    return datetime.now().strftime("%H:%M:%S")

@tool
def get_weather(city: str):
    """Returns the current weather for a specific city."""
    return f"The weather in {city} is sunny and 25°C."

async def run_example_async_agent(bq_project_id: str):
    """Runs an asynchronous Agent with BigQuery logging."""

    # Setup BigQuery logging details
    dataset_id = "your_dataset_id"
    table_id = "your_table_id"

    print(f"--- Starting Async Agent Example ---")
    print(f"Logging to: {bq_project_id}.{dataset_id}.{table_id}")

    # Initialize the async callback handler
    bigquery_handler = AsyncBigQueryCallbackHandler(
        project_id=bq_project_id, dataset_id=dataset_id, table_id=table_id
    )

    try:
        # Setup LLM and Tools
        llm = ChatGoogleGenerativeAI(model="gemini-2.5-pro")
        tools = [get_current_time, get_weather]

        # Create the Agent (LangGraph ReAct implementation)
        agent_executor = create_react_agent(llm, tools)

        query = "What time is it now, and what is the weather like in New York?"
        print(f"User Query: {query}")

        # Run the agent asynchronously
        # We use astream to process chunks, but you can also use ainvoke
        async for chunk in agent_executor.astream(
            {"messages": [("user", query)]},
            config={
                "callbacks": [bigquery_handler],  # Pass the handler here
                "metadata": {
                    "session_id": "agent-session-001",
                    "user_id": "user-123",
                    "agent": "weather-agent",
                },
            },
        ):
            # Print agent thoughts/actions as they happen
            if "agent" in chunk:
                print(f"🤖 Agent: {chunk['agent']['messages'][0].content}")
            elif "tools" in chunk:
                print(f"🛠️ Tool Output: {chunk['tools']['messages'][0].content}")

        print("✅ Agent finished. Logs are being written in the background...")
    finally:
        # Ensure resources are cleaned up
        await bigquery_handler.close()

if __name__ == "__main__":
    # Ensure GOOGLE_API_KEY is set in your environment
    if "GOOGLE_API_KEY" not in os.environ:
        raise ValueError("Please set the GOOGLE_API_KEY environment variable.")

    project_id = "your-project-id"
    asyncio.run(run_example_async_agent(project_id))

import json
import re
from typing import Any

from langchain_google_community.callbacks.bigquery_callback import (
    BigQueryCallbackHandler,
    BigQueryLoggerConfig
)

def redact_dollar_amounts(event_content: Any) -> str:
    """
    Custom formatter to redact dollar amounts (e.g., $600, $12.50)
    and ensure JSON output if the input is a dict.
    """
    text_content = ""
    # If the content is a dictionary (e.g., a list of messages), convert it to a JSON string first.
    if isinstance(event_content, dict):
        text_content = json.dumps(event_content)
    else:
        text_content = str(event_content)

    # Regex to find dollar amounts: $ followed by digits, optionally with commas or decimals.
    # Examples: $600, $1,200.50, $0.99
    redacted_content = re.sub(r'\$\d+(?:,\d{3})*(?:\.\d+)?', 'xxx', text_content)

    return redacted_content

# 1. Configure BigQueryLoggerConfig
config = BigQueryLoggerConfig(
    enabled=True,
    event_allowlist=["LLM_REQUEST", "LLM_RESPONSE"], # Only log these specific events
    shutdown_timeout=10.0,   # Wait up to 10s for logs to flush on exit
    max_content_length=500,  # Truncate content to 500 characters
    content_formatter=redact_dollar_amounts, # Set the custom formatting function
)

# 2. Initialize the Callback Handler
handler = BigQueryCallbackHandler(
    project_id="your-project-id",
    dataset_id="your_dataset",
    table_id="your_table",
    config=config
)

CREATE TABLE `your-gcp-project-id.adk_agent_logs.agent_events_v2`
(
  timestamp TIMESTAMP NOT NULL OPTIONS(description="The UTC timestamp when the event occurred."),
  event_type STRING OPTIONS(description="The category of the event."),
  agent STRING OPTIONS(description="The name of the agent."),
  session_id STRING OPTIONS(description="A unique identifier for the conversation session."),
  invocation_id STRING OPTIONS(description="A unique identifier for a single turn."),
  user_id STRING OPTIONS(description="The identifier of the end-user."),
  trace_id STRING OPTIONS(description="OpenTelemetry trace ID."),
  span_id STRING OPTIONS(description="OpenTelemetry span ID."),
  parent_span_id STRING OPTIONS(description="OpenTelemetry parent span ID."),
  content JSON OPTIONS(description="The primary payload of the event."),
  content_parts ARRAY<STRUCT<
    mime_type STRING,
    uri STRING,
    object_ref STRUCT<
      uri STRING,
      version STRING,
      authorizer STRING,
      details JSON
    >,
    text STRING,
    part_index INT64,
    part_attributes STRING,
    storage_mode STRING
  >> OPTIONS(description="For multi-modal events, contains a list of content parts."),
  attributes JSON OPTIONS(description="Arbitrary key-value pairs."),
  latency_ms JSON OPTIONS(description="Latency measurements."),
  status STRING OPTIONS(description="The outcome of the event."),
  error_message STRING OPTIONS(description="Detailed error message."),
  is_truncated BOOLEAN OPTIONS(description="Flag indicating if content was truncated.")
)
PARTITION BY DATE(timestamp)
CLUSTER BY event_type, agent, user_id;

SELECT
  timestamp,
  event_type,
  span_id,
  parent_span_id,
  -- Extract summary or specific content based on event type
  COALESCE(
    JSON_VALUE(content, '$.messages[0].content'),
    JSON_VALUE(content, '$.summary'),
    JSON_VALUE(content)
  ) AS summary,
  JSON_VALUE(latency_ms, '$.total_ms') AS duration_ms
FROM
  `your-gcp-project-id.adk_agent_logs.agent_events_v2`
WHERE
    -- Replace with a specific trace_id from your logs
  trace_id = '019bb986-a0db-7da1-802d-2725795ab340'
ORDER BY
  timestamp ASC;

SELECT
  JSON_VALUE(attributes, '$.model') AS model,
  COUNT(*) AS total_calls,
  AVG(CAST(JSON_VALUE(latency_ms, '$.total_ms') AS FLOAT64)) AS avg_latency_ms,
  SUM(CAST(JSON_VALUE(attributes, '$.usage.total_tokens') AS INT64)) AS total_tokens
FROM
  `your-gcp-project-id.adk_agent_logs.agent_events_v2`
WHERE
  event_type = 'LLM_RESPONSE'
GROUP BY
  1;

SELECT
  logs.session_id,
  -- Get a signed URL for the image (optional, for viewing)
  STRING(OBJ.GET_ACCESS_URL(parts.object_ref, "r").access_urls.read_url) as signed_url,
  -- Analyze the image using a remote model (e.g., gemini-1.5-pro)
  AI.GENERATE(
    ('Describe this image briefly. What company logo?', parts.object_ref)
  ) AS generated_result
FROM
  `your-gcp-project-id.adk_agent_logs.agent_events_v2` logs,
  UNNEST(logs.content_parts) AS parts
WHERE
  parts.mime_type LIKE 'image/%'
ORDER BY logs.timestamp DESC
LIMIT 1;

SELECT
  timestamp,
  event_type,
  part.mime_type,
  part.storage_mode,
  part.object_ref.uri AS gcs_uri,
  -- Generate a signed URL to read the content directly (requires connection_id configuration)
  STRING(OBJ.GET_ACCESS_URL(part.object_ref, 'r').access_urls.read_url) AS signed_url
FROM `your-gcp-project-id.adk_agent_logs.agent_events_v2`,
UNNEST(content_parts) AS part
WHERE part.storage_mode = 'GCS_REFERENCE'
ORDER BY timestamp DESC
LIMIT 10;