Skip to content

BigQueryAgentAnalyticsPlugin: Pass **kwargs to _log_event for all callbacks (not just model callbacks) #4330

@maryhamilton

Description

@maryhamilton

** Please make sure you read the contribution guide and file the issues in the right place. **
Contribution guide.

🔴 Required Information

Please ensure all items in this section are completed to allow for efficient
triaging. Requests without complete information may be rejected / deprioritized.
If an item is not applicable to you - please mark it as N/A

Is your feature request related to a specific problem?

Please describe the problem you are trying to solve. (Ex: "I'm always frustrated
when I have to manually handle X...")

The BigQueryAgentAnalyticsPlugin callback methods inconsistently handle **kwargs. Model callbacks (before_model_callback, after_model_callback) properly pass **kwargs through to _log_event, which serializes them into the BigQuery attributes JSON field. However, agent callbacks, tool callbacks, and run callbacks accept **kwargs but do not pass them to _log_event, effectively dropping any custom attributes.
This makes it impossible to add custom metadata (like customer_id, cost_center, email_address) to non-model events through the documented callback override pattern.

Describe the Solution You'd Like

A clear and concise description of the feature or API change you want.
Be specific about input/outputs if this involves an API change.

All callback methods that accept **kwargs should pass them through to _log_event so custom attributes are serialized into the attributes JSON field for ALL event types.
Affected methods that currently drop **kwargs:

  • before_agent_callback / after_agent_callback
  • before_tool_callback / after_tool_callback
  • before_run_callback / after_run_callback
  • on_user_message_callback
  • on_tool_error_callback

Impact on your work

How does this feature impact your work and what are you trying to achieve?
If this is critical for you, tell us if there is a timeline by when you need
this feature.

We need to track customer_id across all BigQuery analytics events for multi-tenant reporting. Currently, customer_id only appears in LLM_REQUEST and LLM_RESPONSE events, but is missing from AGENT_STARTING, AGENT_COMPLETED, TOOL_STARTING, TOOL_COMPLETED, and INVOCATION_STARTING/COMPLETED events.
This is important for production use but we have a workaround in place.

Willingness to contribute

Are you interested in implementing this feature yourself or submitting a PR?
(Yes/No)

No


🟡 Recommended Information

Describe Alternatives You've Considered

A clear and concise description of any alternative solutions or workarounds
you've considered and why they didn't work for you.

Workaround implemented: We override the internal _log_event method to inject customer_id:

class CustomerAwareBqAnalyticsPlugin(BigQueryAgentAnalyticsPlugin):
    async def _log_event(
        self,
        event_type: str,
        callback_context: CallbackContext,
        raw_content: Any = None,
        is_truncated: bool = False,
        **kwargs: Any,
    ) -> None:
        customer_id = callback_context.state.get("customer_id")
        if customer_id is not None:
            kwargs["customer_id"] = customer_id
        await super()._log_event(event_type, callback_context, raw_content, is_truncated, **kwargs)

This works but is fragile because _log_event is an internal method that may change between versions.

Proposed API / Implementation

If you have ideas on how this should look in code, please share a
pseudo-code example.

For each affected callback, pass **kwargs to _log_event. Example fix for before_agent_callback:

# Current (drops kwargs):
async def before_agent_callback(
    self, *, agent: Any, callback_context: CallbackContext, **kwargs
) -> None:
    TraceManager.init_trace(callback_context)
    TraceManager.push_span(callback_context, "agent")
    await self._log_event(
        "AGENT_STARTING",
        callback_context,
        raw_content=getattr(agent, "instruction", ""),
    )

# Proposed (passes kwargs):
async def before_agent_callback(
    self, *, agent: Any, callback_context: CallbackContext, **kwargs
) -> None:
    TraceManager.init_trace(callback_context)
    TraceManager.push_span(callback_context, "agent")
    await self._log_event(
        "AGENT_STARTING",
        callback_context,
        raw_content=getattr(agent, "instruction", ""),
        **kwargs,  # <-- Add this
    )

Additional Context

Add any other context or screenshots about the feature request here.

The documentation in bigquery_agent_analytics_plugin.py states that any additional keyword arguments should be serialized into the attributes field, but this only works for model callbacks. Making this consistent across all callbacks would match the expected behaviour.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bq[Component] This issue is related to Big Query integration

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions