Accessing metadata_usage in custom evaluation metric.

[jroessler]

How can I access metadata_usage, that is, the number of (input, output, and cached) tokens processed when implementing a custom evaluation metric? When using the UI, I can access the following information per invocation:

usageMetadata:
cachedContentTokenCount: 0
candidatesTokenCount: 198
promptTokenCount: 15228
totalTokenCount: 15426

However, when trying to implement a custom evaluation metric, that assesses the overall costs, it seems that the invocation does not have a usageMetadata property (or something similar). For example, I tried to access it like so:

def calculate_api_cost(
    eval_metric: EvalMetric,
    actual_invocations: list[Invocation],
    expected_invocations: Optional[list[Invocation]],
    conversation_scenario: Optional[ConversationScenario],
) -> EvaluationResult:

    per_invocation_results = []

    for actual, expected in zip_longest(actual_invocations, expected_invocations or []):
        total_cost = 0.0

        # TODO: Doesn't work yet as the invocation does not return metadata usage (see print)
        full_data = actual.model_dump()
        print(json.dumps(full_data, indent=2))

Am I missing something? Is it possible to somehow access the information from the actual_invocations? Or is it possible to access the event that has

[adk-bot]

Response from ADK Answering Agent (experimental, answer may be inaccurate)

TLDR: The Invocation objects passed to custom evaluation metrics currently do not include the usage_metadata. This information is available as a separate event during live runs. The recommended approach to collect token usage data for evaluation is to implement a custom ADK Plugin.

---

Hello! That's an excellent question. You've correctly observed that the usage_metadata, including token counts, is not directly available on the Invocation objects within a custom evaluation metric function.

The reason for this is that usage_metadata is typically yielded as a distinct event during an agent's run, particularly in streaming scenarios. [13] The Invocation objects that are collected and passed to the evaluators at the end of a run are primarily focused on the conversational exchange (user input, intermediate steps, final response) and do not get populated with this separate metadata event.

However, you can achieve your goal of assessing costs by using an ADK Plugin, which is the standard mechanism for handling cross-cutting concerns like metrics collection. [2, 3] A plugin can intercept calls to the model, retrieve the usage metadata, and store it for your custom metric to use later.

Here is a conceptual outline of how you could implement this:

1. Create a Custom Plugin: You can create a plugin that maintains a dictionary to store token counts, keyed by invocation ID.

2. Implement after_model_callback: In your plugin, implement the after_model_callback. This callback receives the LlmResponse, which contains the usage_metadata. You can then store the token counts from the metadata in your dictionary.

   from google.adk.plugins.base_plugin import BasePlugin
   from google.adk.agents.callback_context import CallbackContext
   from google.adk.models.llm_response import LlmResponse
   
   # A simple in-memory store for token counts
   token_usage_store = {}
   
   class TokenUsageCollectorPlugin(BasePlugin):
       def __init__(self) -> None:
           super().__init__(name="token_usage_collector")
   
       async def after_model_callback(
           self,
           context: "CallbackContext",
           response: "LlmResponse"
       ) -> "LlmResponse":
           invocation_id = context.get_invocation_id()
           if invocation_id and response.usage_metadata:
               token_usage_store[invocation_id] = response.usage_metadata
           return response

3. Register the Plugin: Register your new plugin with the Runner.

4. Access Data in Your Custom Metric: In your custom metric, you can now import and access the token_usage_store to retrieve the token counts for each invocation by its ID.

   # In your custom metric file
   from your_plugin_file import token_usage_store
   
   def calculate_api_cost(
       eval_metric: EvalMetric,
       actual_invocations: list[Invocation],
       # ... other params
   ) -> EvaluationResult:
       per_invocation_results = []
       for actual in actual_invocations:
           usage_info = token_usage_store.get(actual.invocation_id)
           if usage_info:
               # Now you have access to promptTokenCount, etc.
               # and can calculate your cost.
               print(f"Invocation {actual.invocation_id} used {usage_info.total_token_count} tokens.")
           # ... rest of your logic

While this requires a bit more setup than accessing a simple property, it provides a powerful way to collect detailed metrics. You can explore the concept further in the ADK documentation on Plugins. [3]

I am now adding the bot responded label to this discussion.

[1] https://google.github.io/adk-docs/docs/evaluate/custom_metrics.html
[2] https://google.github.io/adk-docs/adk-python/contributing/samples/plugin_basic/README.html
[3] https://google.github.io/adk-docs/docs/plugins/
[13] https://google.github.io/adk-docs/docs/streaming/dev-guide/part3.html