gen_ai.client.token.usage records nothing when the last response of a model call has no usage #7560

import asyncio
import os

os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")

from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.run_config import RunConfig
from google.adk.agents.run_config import StreamingMode
from google.adk.models.lite_llm import LiteLlm
from google.adk.models.lite_llm import LiteLLMClient
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
import google.adk.telemetry._metrics as adk_metrics
from google.genai import types
from litellm import Function
from litellm.types.utils import ChatCompletionDeltaToolCall
from litellm.types.utils import Delta
from litellm.types.utils import ModelResponseStream
from litellm.types.utils import StreamingChoices
from opentelemetry import metrics
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import InMemoryMetricReader

reader = InMemoryMetricReader()
metrics.set_meter_provider(MeterProvider(metric_readers=[reader]))


class FakeClient(LiteLLMClient):
  """Replays a stream that stops at max_tokens inside a tool call."""

  async def acompletion(self, model, messages, tools, **kwargs):
    async def stream():
      yield ModelResponseStream(
          choices=[
              StreamingChoices(
                  finish_reason="length",
                  delta=Delta(
                      role="assistant",
                      content="Let me check the weather.",
                      tool_calls=[
                          ChatCompletionDeltaToolCall(
                              type="function",
                              id="call_1",
                              index=0,
                              function=Function(
                                  name="get_weather",
                                  arguments='{"city": "Par',
                              ),
                          )
                      ],
                  ),
              )
          ]
      )
      yield ModelResponseStream(
          usage={
              "prompt_tokens": 120,
              "completion_tokens": 40,
              "total_tokens": 160,
          },
          choices=[StreamingChoices(finish_reason=None)],
      )

    return stream()


def get_weather(city: str) -> str:
  return f"Sunny in {city}"


def token_points(name):
  points = {}
  data = reader.get_metrics_data()
  for rm in data.resource_metrics if data else []:
    for sm in rm.scope_metrics:
      for metric in sm.metrics:
        if metric.name == name:
          for point in metric.data.data_points:
            key = dict(point.attributes).get("gen_ai.token.type")
            points[key] = point.sum
  return points


async def main():
  print("adk module:", adk_metrics.__file__)
  agent = LlmAgent(
      name="weather",
      model=LiteLlm(model="openai/gpt-4o", llm_client=FakeClient()),
      tools=[get_weather],
  )
  runner = Runner(
      app_name="app", agent=agent, session_service=InMemorySessionService()
  )
  await runner.session_service.create_session(
      app_name="app", user_id="u", session_id="s"
  )
  async for event in runner.run_async(
      user_id="u",
      session_id="s",
      new_message=types.UserContent("Weather in Paris?"),
      run_config=RunConfig(streaming_mode=StreamingMode.SSE),
  ):
    if not event.partial:
      usage = event.usage_metadata
      print(
          "event:",
          event.finish_reason,
          "usage:",
          (usage.prompt_token_count, usage.candidates_token_count)
          if usage
          else None,
      )
  print("gen_ai.client.token.usage:", token_points("gen_ai.client.token.usage"))


asyncio.run(main())