-
Notifications
You must be signed in to change notification settings - Fork 2.8k
Description
** 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_callbackbefore_tool_callback/after_tool_callbackbefore_run_callback/after_run_callbackon_user_message_callbackon_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.