diff --git a/src/google/adk/events/event.py b/src/google/adk/events/event.py index 2c6a6cd66c..0383e41b25 100644 --- a/src/google/adk/events/event.py +++ b/src/google/adk/events/event.py @@ -66,6 +66,19 @@ class Event(LlmResponse): conversation history. """ + tool_usage_metadata: Optional[dict[str, types.GenerateContentResponseUsageMetadata]] = None + """Token usage metadata for tools and sub-agents invoked during this event. + + Maps tool/agent names to their respective usage metadata, enabling granular + cost tracking for nested agent architectures and Vertex AI features. + + Example: + { + "vertex_ai_search": UsageMetadata(prompt_token_count=100, ...), + "sub_agent_name": UsageMetadata(prompt_token_count=500, ...), + } + """ + # The following are computed fields. # Do not assign the ID. It will be assigned by the session. id: str = '' diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 664d0d6c67..6aac8eb6aa 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -962,12 +962,16 @@ def __build_response_event( parts=[part_function_response], ) + # Collect tool usage metadata from tool context + tool_usage = tool_context.get_all_tool_usage() + function_response_event = Event( invocation_id=invocation_context.invocation_id, author=invocation_context.agent.name, content=content, actions=tool_context.actions, branch=invocation_context.branch, + tool_usage_metadata=tool_usage if tool_usage else None, ) return function_response_event diff --git a/src/google/adk/models/llm_response.py b/src/google/adk/models/llm_response.py index 754e5abcfb..32e785ce5b 100644 --- a/src/google/adk/models/llm_response.py +++ b/src/google/adk/models/llm_response.py @@ -142,6 +142,48 @@ class LlmResponse(BaseModel): It can be used to identify and chain interactions for stateful conversations. """ + @staticmethod + def merge_usage_metadata( + metadata_list: list[Optional[types.GenerateContentResponseUsageMetadata]] + ) -> Optional[types.GenerateContentResponseUsageMetadata]: + """Merges multiple usage metadata objects into a single aggregate. + + Args: + metadata_list: List of usage metadata objects to merge. + + Returns: + Merged usage metadata with cumulative token counts, or None if all inputs + are None. + """ + if not metadata_list or all(m is None for m in metadata_list): + return None + + total_prompt_tokens = 0 + total_candidates_tokens = 0 + total_tokens = 0 + total_cached_tokens = 0 + + for metadata in metadata_list: + if metadata: + total_prompt_tokens += metadata.prompt_token_count or 0 + total_candidates_tokens += metadata.candidates_token_count or 0 + total_tokens += metadata.total_token_count or 0 + if hasattr(metadata, 'cached_content_token_count'): + total_cached_tokens += metadata.cached_content_token_count or 0 + + # Create merged metadata + merged = types.GenerateContentResponseUsageMetadata( + prompt_token_count=total_prompt_tokens, + candidates_token_count=total_candidates_tokens, + total_token_count=total_tokens, + ) + + # Add cached tokens if any were present + if total_cached_tokens > 0: + merged.cached_content_token_count = total_cached_tokens + + return merged + @staticmethod def create( generate_content_response: types.GenerateContentResponse, diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 7cbf931ca9..981e8d9281 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -2364,6 +2364,24 @@ async def after_model_callback( if usage_dict: content_dict["usage"] = usage_dict + # Add tool-level usage metadata from nested agents and Vertex AI features + if hasattr(llm_response, "tool_usage_metadata") and llm_response.tool_usage_metadata: + tool_usage_dict = {} + for tool_name, tool_usage in llm_response.tool_usage_metadata.items(): + if tool_usage: + tool_usage_entry = {} + if hasattr(tool_usage, "prompt_token_count"): + tool_usage_entry["prompt"] = tool_usage.prompt_token_count + if hasattr(tool_usage, "candidates_token_count"): + tool_usage_entry["completion"] = tool_usage.candidates_token_count + if hasattr(tool_usage, "total_token_count"): + tool_usage_entry["total"] = tool_usage.total_token_count + if tool_usage_entry: + tool_usage_dict[tool_name] = tool_usage_entry + + if tool_usage_dict: + content_dict["tool_usage"] = tool_usage_dict + if content_dict: content_str = content_dict else: diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py index cd6b7071b8..57bc7d5ade 100644 --- a/src/google/adk/telemetry/tracing.py +++ b/src/google/adk/telemetry/tracing.py @@ -214,6 +214,39 @@ def trace_tool_call( else: span.set_attribute('gcp.vertex.agent.tool_response', '{}') + # Add tool-level usage metadata if available + if ( + function_response_event is not None + and function_response_event.tool_usage_metadata + ): + total_prompt_tokens = 0 + total_completion_tokens = 0 + total_tokens = 0 + + for tool_name, usage_metadata in function_response_event.tool_usage_metadata.items(): + if usage_metadata: + total_prompt_tokens += getattr(usage_metadata, 'prompt_token_count', 0) or 0 + total_completion_tokens += getattr(usage_metadata, 'candidates_token_count', 0) or 0 + total_tokens += getattr(usage_metadata, 'total_token_count', 0) or 0 + + if total_tokens > 0: + span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, total_prompt_tokens) + span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, total_completion_tokens) + span.set_attribute('gcp.vertex.agent.tool_usage_total_tokens', total_tokens) + + # Add detailed breakdown as custom attribute + span.set_attribute( + 'gcp.vertex.agent.tool_usage_breakdown', + _safe_json_serialize({ + name: { + 'prompt_tokens': getattr(usage, 'prompt_token_count', 0) or 0, + 'completion_tokens': getattr(usage, 'candidates_token_count', 0) or 0, + 'total_tokens': getattr(usage, 'total_token_count', 0) or 0, + } + for name, usage in function_response_event.tool_usage_metadata.items() + }) + ) + def trace_merged_tool_calls( response_event_id: str, diff --git a/src/google/adk/tools/agent_tool.py b/src/google/adk/tools/agent_tool.py index 91135dce5f..3d11d9aa4b 100644 --- a/src/google/adk/tools/agent_tool.py +++ b/src/google/adk/tools/agent_tool.py @@ -244,6 +244,8 @@ async def run_async( state=state_dict, ) + # Aggregate usage metadata from sub-agent execution + usage_metadata_list = [] last_content = None async with Aclosing( runner.run_async( @@ -251,6 +253,15 @@ async def run_async( ) ) as agen: async for event in agen: + # Collect usage metadata from each event + if event.usage_metadata: + usage_metadata_list.append(event.usage_metadata) + + # Also collect tool-level usage from nested events + if event.tool_usage_metadata: + for tool_name, tool_usage in event.tool_usage_metadata.items(): + usage_metadata_list.append(tool_usage) + # Forward state delta to parent session. if event.actions.state_delta: tool_context.state.update(event.actions.state_delta) @@ -261,6 +272,13 @@ async def run_async( # to avoid "Attempted to exit cancel scope in a different task" errors await runner.close() + # Aggregate and record usage for this sub-agent + if usage_metadata_list: + from ..models.llm_response import LlmResponse + aggregated_usage = LlmResponse.merge_usage_metadata(usage_metadata_list) + if aggregated_usage: + tool_context.set_tool_usage(self.agent.name, aggregated_usage) + if last_content is None or last_content.parts is None: return '' merged_text = '\n'.join( diff --git a/src/google/adk/tools/tool_context.py b/src/google/adk/tools/tool_context.py index 12ae55dbc9..ed353a854c 100644 --- a/src/google/adk/tools/tool_context.py +++ b/src/google/adk/tools/tool_context.py @@ -59,11 +59,45 @@ def __init__( super().__init__(invocation_context, event_actions=event_actions) self.function_call_id = function_call_id self.tool_confirmation = tool_confirmation + self._tool_usage: dict[str, Any] = {} @property def actions(self) -> EventActions: return self._event_actions + def set_tool_usage( + self, + tool_name: str, + usage_metadata: Any, + ) -> None: + """Records usage metadata for a tool or sub-agent invocation. + + Args: + tool_name: Name of the tool or agent that generated usage. + usage_metadata: Usage metadata object (GenerateContentResponseUsageMetadata + or dict with token counts). + """ + self._tool_usage[tool_name] = usage_metadata + + def get_tool_usage(self, tool_name: str) -> Optional[Any]: + """Retrieves usage metadata for a specific tool. + + Args: + tool_name: Name of the tool to retrieve usage for. + + Returns: + Usage metadata if recorded, None otherwise. + """ + return self._tool_usage.get(tool_name) + + def get_all_tool_usage(self) -> dict[str, Any]: + """Returns all tool usage metadata recorded in this context. + + Returns: + Dictionary mapping tool names to their usage metadata. + """ + return self._tool_usage.copy() + def request_credential(self, auth_config: AuthConfig) -> None: if not self.function_call_id: raise ValueError('function_call_id is not set.')