Skip to content

feat(google-genai): extract modality token usage attributes on generate_content #674

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lmolkova merged 8 commits into from
Sep 16, 2026
Merged

feat(google-genai): extract modality token usage attributes on generate_content #674

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
cb50b61
feat(google-genai): extract modality token usage attributes on genera…
Krishnachaitanyakc Sep 10, 2026
0530570
Add changelog fragment
Krishnachaitanyakc Sep 10, 2026
1680397
Merge branch 'main' into google-genai-modality-token-usage
Krishnachaitanyakc Sep 10, 2026
d1a3a1a
refactor(util-genai): own the modality token mapping on InferenceInvo…
Krishnachaitanyakc Sep 11, 2026
72c44f9
Merge branch 'main' into google-genai-modality-token-usage
Krishnachaitanyakc Sep 11, 2026
12ba4b1
Avoid walking property paths on every intermediate streaming chunk th…
Krishnachaitanyakc Sep 11, 2026
7dd8e9c
Merge remote-tracking branch 'origin/main' into google-genai-modality…
Krishnachaitanyakc Sep 16, 2026
2c9dadb
Apply ruff format
Krishnachaitanyakc Sep 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions instrumentation/opentelemetry-instrumentation-google-genai/.changelog/674.added
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Record modality token usage breakdown attributes (text, image, audio) for generate_content.
68 changes: 67 additions & 1 deletion ...mentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

import json
import os
from collections.abc import AsyncIterable, Callable, Iterable
from collections.abc import AsyncIterable, Callable, Iterable, Mapping
from typing import Any

from google.genai.models import AsyncModels, Models
Expand All @@ -14,6 +14,7 @@
GenerateContentConfig,
GenerateContentConfigOrDict,
GenerateContentResponse,
ModalityTokenCount,
Tool,
ToolUnionDict,
)
Expand Down Expand Up @@ -346,6 +347,50 @@ def _apply_request_attributes(
)


_INPUT_MODALITY_FIELDS = {
"text": "text_input_tokens",
"image": "image_input_tokens",
"audio": "audio_input_tokens",
}
_OUTPUT_MODALITY_FIELDS = {
"text": "text_output_tokens",
"image": "image_output_tokens",
"audio": "audio_output_tokens",
}
_CACHE_READ_MODALITY_FIELDS = {
"text": "text_cache_read_input_tokens",
"image": "image_cache_read_input_tokens",
"audio": "audio_cache_read_input_tokens",
}


def _set_modality_tokens(
invocation: InferenceInvocation,
entries: list[ModalityTokenCount] | None,
fields: Mapping[str, str],
) -> None:
if entries is None:
return
# A streaming chunk redelivers the whole breakdown, so an arriving one has
# to replace the previous values: a modality that drops out of a later
# chunk would otherwise keep a stale count and push the breakdown over the
# total reported alongside it.
for field_name in fields.values():
setattr(invocation, field_name, None)
for entry in entries:
modality = entry.modality
token_count = entry.token_count
if modality is None or not isinstance(token_count, int):
continue
# modality is a MediaModality enum, so str() would yield
# "MediaModality.AUDIO" rather than the bare modality name.
field_name = fields.get(
str(getattr(modality, "value", modality)).lower()
)
if field_name is not None:
setattr(invocation, field_name, token_count)


@lmolkova lmolkova Sep 11, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the third copy of this mapping - interactions.py has one and langchain got one in #671. Let's move this onto InferenceInvocation in opentelemetry-util-genai instead of duplicating and trying to unify later.

Rough sketch (not final, needs details sorted out):

class InferenceInvocation(...):
    def set_input_tokens(self, modality: Modality | str, tokens: int | None) -> None:
        self._set_modality_tokens(_INPUT, modality, tokens)

    def set_output_tokens(...) -> None: ...
    def set_cache_read_input_tokens(...) -> None: ...

    def _set_modality_tokens(self, bucket, modality, tokens) -> None:
        if not isinstance(tokens, int) or isinstance(tokens, bool) or tokens < 0:
            return
        name = str(getattr(modality, "value", modality)).lower()
        # drop unsupported modalities (video, doc) centrally
        ...

Call site here simplifies to:

for entry in _get_response_property(response, "usage_metadata.prompt_tokens_details") or []:
    invocation.set_input_tokens(entry.modality, entry.token_count)

This removes the local dicts, enum/string normalization, and int checks. It also centralizes handling unmapped modalities and future token types across all instrumentations.

Note: Modality in types.py currently lacks text, need to add it.

It's a small change - let's add the util method and update interactions.py and langchain right away

@Krishnachaitanyakc Krishnachaitanyakc Sep 11, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, pushed in d1a3a1a. Went with the whole iterable rather than per entry, so an arriving breakdown replaces rather than merges. Happy to switch.

interactions.py is in there, which #616 also touches.

def _get_response_property(response: GenerateContentResponse, path: str):
path_segments = path.split(".")
current_context = response
Expand Down Expand Up @@ -423,6 +468,27 @@ def _apply_response_attributes(
invocation.output_tokens = (
invocation.output_tokens or 0
) + thinking_tokens
_set_modality_tokens(
invocation,
_get_response_property(
response, "usage_metadata.prompt_tokens_details"
),
Comment thread
Krishnachaitanyakc marked this conversation as resolved.
Outdated
_INPUT_MODALITY_FIELDS,
)
_set_modality_tokens(
invocation,
_get_response_property(
response, "usage_metadata.candidates_tokens_details"
),
_OUTPUT_MODALITY_FIELDS,
)
_set_modality_tokens(
invocation,
_get_response_property(
response, "usage_metadata.cache_tokens_details"
),
_CACHE_READ_MODALITY_FIELDS,
)


def _maybe_get_tool_definitions(
Expand Down
86 changes: 86 additions & 0 deletions ...on/opentelemetry-instrumentation-google-genai/tests/generate_content/nonstreaming_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,10 @@
from google.genai.types import (
FunctionDeclarationDict,
GenerateContentConfig,
GenerateContentResponseUsageMetadata,
GoogleMaps,
MediaModality,
ModalityTokenCount,
ToolDict,
)
from pydantic import BaseModel, Field
Expand Down Expand Up @@ -337,6 +340,89 @@ def test_generated_span_counts_tokens(self):
span.attributes["gen_ai.usage.reasoning.output_tokens"], 17
)

def test_generated_span_counts_modality_tokens(self):
usage_metadata = GenerateContentResponseUsageMetadata(
prompt_tokens_details=[
ModalityTokenCount(
modality=MediaModality.TEXT, token_count=11
),
ModalityTokenCount(
modality=MediaModality.AUDIO, token_count=22
),
],
candidates_tokens_details=[
ModalityTokenCount(
modality=MediaModality.AUDIO, token_count=33
),
ModalityTokenCount(
modality=MediaModality.IMAGE, token_count=44
),
],
cache_tokens_details=[
ModalityTokenCount(
modality=MediaModality.TEXT, token_count=55
),
ModalityTokenCount(
modality=MediaModality.IMAGE, token_count=66
),
],
)
self.configure_valid_response(usage_metadata=usage_metadata)
self.generate_content(model="gemini-2.0-flash", contents="Some input")
span = self.otel.get_span_named("generate_content gemini-2.0-flash")
for attribute, value in (
("gen_ai.usage.text.input_tokens", 11),
("gen_ai.usage.audio.input_tokens", 22),
("gen_ai.usage.audio.output_tokens", 33),
("gen_ai.usage.image.output_tokens", 44),
("gen_ai.usage.text.cache_read.input_tokens", 55),
("gen_ai.usage.image.cache_read.input_tokens", 66),
):
self.assertIsInstance(span.attributes[attribute], int)
self.assertEqual(span.attributes[attribute], value)
for attribute in (
"gen_ai.usage.image.input_tokens",
"gen_ai.usage.text.output_tokens",
"gen_ai.usage.audio.cache_read.input_tokens",
):
self.assertNotIn(attribute, span.attributes)

def test_generated_span_skips_unmapped_modality_tokens(self):
usage_metadata = GenerateContentResponseUsageMetadata(
prompt_tokens_details=[
ModalityTokenCount(
modality=MediaModality.VIDEO, token_count=7
),
ModalityTokenCount(modality=MediaModality.TEXT),
],
)
self.configure_valid_response(usage_metadata=usage_metadata)
self.generate_content(model="gemini-2.0-flash", contents="Some input")
span = self.otel.get_span_named("generate_content gemini-2.0-flash")
for attribute in (
"gen_ai.usage.text.input_tokens",
"gen_ai.usage.image.input_tokens",
"gen_ai.usage.audio.input_tokens",
):
self.assertNotIn(attribute, span.attributes)

def test_generated_span_omits_modality_tokens_when_absent(self):
self.configure_valid_response(input_tokens=1, output_tokens=2)
self.generate_content(model="gemini-2.0-flash", contents="Some input")
span = self.otel.get_span_named("generate_content gemini-2.0-flash")
for attribute in (
"gen_ai.usage.text.input_tokens",
"gen_ai.usage.image.input_tokens",
"gen_ai.usage.audio.input_tokens",
"gen_ai.usage.text.output_tokens",
"gen_ai.usage.image.output_tokens",
"gen_ai.usage.audio.output_tokens",
"gen_ai.usage.text.cache_read.input_tokens",
"gen_ai.usage.image.cache_read.input_tokens",
"gen_ai.usage.audio.cache_read.input_tokens",
):
self.assertNotIn(attribute, span.attributes)

def test_generated_span_records_response_model(self):
self.configure_valid_response(model_version="gemini-2.0-flash-001")
self.generate_content(model="gemini-2.0-flash", contents="Some input")
Expand Down
62 changes: 62 additions & 0 deletions ...ation/opentelemetry-instrumentation-google-genai/tests/generate_content/streaming_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,12 @@

import unittest

from google.genai.types import (
GenerateContentResponseUsageMetadata,
MediaModality,
ModalityTokenCount,
)

from opentelemetry import context as context_api
from opentelemetry.instrumentation.google_genai import (
GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY,
Expand Down Expand Up @@ -93,6 +99,62 @@ def test_includes_token_counts_in_span_not_aggregated_from_responses(self):
self.assertEqual(span.attributes["gen_ai.usage.input_tokens"], 3)
self.assertEqual(span.attributes["gen_ai.usage.output_tokens"], 5)

def test_modality_token_counts_are_replaced_not_merged(self):
# Cumulative counts are returned on each response, so the last
# breakdown wins outright and a modality that drops out of it must not
# survive from an earlier chunk.
self.configure_valid_response(
usage_metadata=GenerateContentResponseUsageMetadata(
prompt_tokens_details=[
ModalityTokenCount(
modality=MediaModality.TEXT, token_count=10
),
ModalityTokenCount(
modality=MediaModality.AUDIO, token_count=90
),
],
),
response_id="qwerty17",
)
self.configure_valid_response(
usage_metadata=GenerateContentResponseUsageMetadata(
prompt_tokens_details=[
ModalityTokenCount(
modality=MediaModality.TEXT, token_count=12
),
],
),
response_id="qwerty18",
)

self.generate_content(model="gemini-2.0-flash", contents="Some input")

span = self.otel.get_span_named("generate_content gemini-2.0-flash")
self.assertEqual(span.attributes["gen_ai.response.id"], "qwerty18")
self.assertEqual(span.attributes["gen_ai.usage.text.input_tokens"], 12)
self.assertNotIn("gen_ai.usage.audio.input_tokens", span.attributes)

def test_modality_token_counts_survive_a_chunk_without_a_breakdown(self):
self.configure_valid_response(
usage_metadata=GenerateContentResponseUsageMetadata(
prompt_tokens_details=[
ModalityTokenCount(
modality=MediaModality.AUDIO, token_count=90
),
],
),
response_id="qwerty17",
)
self.configure_valid_response(response_id="qwerty18")

self.generate_content(model="gemini-2.0-flash", contents="Some input")

span = self.otel.get_span_named("generate_content gemini-2.0-flash")
self.assertEqual(span.attributes["gen_ai.response.id"], "qwerty18")
self.assertEqual(
span.attributes["gen_ai.usage.audio.input_tokens"], 90
)

def test_log_has_extra_genai_attributes(self):
self.configure_valid_response(text="Yep, it works!")
tok = context_api.attach(
Expand Down