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

instrumentation/opentelemetry-instrumentation-genai-langchain/.changelog/674.changed

@@ -0,0 +1 @@
+Record modality token usage through the shared `InferenceInvocation` setters; `extract_token_details` no longer returns modality keys.

instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/callback_handler.py

@@ -36,6 +36,7 @@
     is_stream_end_marker,
     make_input_message,
     make_last_output_message,
+    modality_tokens,
     normalize_provider,
     prepare_tool_definitions,
     resolve_response_model_and_id,
@@ -561,35 +562,16 @@ def on_llm_end(
                         ) is not None:
                             llm_invocation.thinking_tokens = reasoning_tokens
 
-                        if (
-                            text_in := token_details.get("text_input_tokens")
-                        ) is not None:
-                            llm_invocation.text_input_tokens = text_in
-                        if (
-                            image_in := token_details.get("image_input_tokens")
-                        ) is not None:
-                            llm_invocation.image_input_tokens = image_in
-                        if (
-                            audio_in := token_details.get("audio_input_tokens")
-                        ) is not None:
-                            llm_invocation.audio_input_tokens = audio_in
-
-                        if (
-                            text_out := token_details.get("text_output_tokens")
-                        ) is not None:
-                            llm_invocation.text_output_tokens = text_out
-                        if (
-                            image_out := token_details.get(
-                                "image_output_tokens"
+                        llm_invocation.set_input_tokens(
+                            modality_tokens(
+                                usage_metadata, "input_token_details"
                             )
-                        ) is not None:
-                            llm_invocation.image_output_tokens = image_out
-                        if (
-                            audio_out := token_details.get(
-                                "audio_output_tokens"
+                        )
+                        llm_invocation.set_output_tokens(
+                            modality_tokens(
+                                usage_metadata, "output_token_details"
                             )
-                        ) is not None:
-                            llm_invocation.audio_output_tokens = audio_out
+                        )
 
                         llm_invocation.output_tokens = output_tokens
 

instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/utils.py

@@ -29,6 +29,7 @@
     FunctionToolDefinition,
     InputMessage,
     MessagePart,
+    ModalityTokens,
     OutputMessage,
     ReasoningPart,
     Role,
@@ -617,7 +618,7 @@ def resolve_response_model_and_id(
 
 
 def extract_token_details(usage_metadata: dict[str, Any]) -> dict[str, int]:
-    """Extract cache, reasoning, and modality token break-downs from LangChain usage metadata."""
+    """Extract cache and reasoning token break-downs from LangChain usage metadata."""
 
     token_details: dict[str, int] = {}
     raw_input_details = usage_metadata.get("input_token_details")
@@ -655,20 +656,19 @@ def _get_positive_int(d: dict[str, Any], key: str) -> int | None:
     ) is not None:
         token_details["reasoning_tokens"] = reasoning
 
-    # Input modality breakdowns
-    if (text_in := _get_positive_int(input_details, "text")) is not None:
-        token_details["text_input_tokens"] = text_in
-    if (image_in := _get_positive_int(input_details, "image")) is not None:
-        token_details["image_input_tokens"] = image_in
-    if (audio_in := _get_positive_int(input_details, "audio")) is not None:
-        token_details["audio_input_tokens"] = audio_in
-
-    # Output modality breakdowns
-    if (text_out := _get_positive_int(output_details, "text")) is not None:
-        token_details["text_output_tokens"] = text_out
-    if (image_out := _get_positive_int(output_details, "image")) is not None:
-        token_details["image_output_tokens"] = image_out
-    if (audio_out := _get_positive_int(output_details, "audio")) is not None:
-        token_details["audio_output_tokens"] = audio_out
-
     return token_details
+
+
+def modality_tokens(
+    usage_metadata: Mapping[str, Any], key: str
+) -> ModalityTokens | None:
+    """Read one of LangChain's ``*_token_details`` maps as modality pairs.
+
+    Returns ``None`` when the key is absent or not a mapping, so the caller
+    leaves any previously recorded breakdown alone. Non-modality keys such as
+    ``cache_read`` and ``reasoning`` are dropped downstream.
+    """
+    details = usage_metadata.get(key)
+    if not isinstance(details, Mapping):
+        return None
+    return list(cast("Mapping[str, int | None]", details).items())

instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_callback_handler.py

@@ -1663,14 +1663,48 @@ def test_modality_tokens_set_on_invocation(self):
 
         handler.on_llm_end(response=response, run_id=run_id)
 
+        # The breakdown is applied by InferenceInvocation, so the handler's
+        # contract is the pairs it forwards. util-genai covers the mapping
+        # from those pairs onto span attributes.
         assert llm_inv.input_tokens == 100
-        assert llm_inv.text_input_tokens == 70
-        assert llm_inv.image_input_tokens == 20
-        assert llm_inv.audio_input_tokens == 10
         assert llm_inv.output_tokens == 50
-        assert llm_inv.text_output_tokens == 35
-        assert llm_inv.image_output_tokens == 10
-        assert llm_inv.audio_output_tokens == 5
+        # LangChain's pydantic model reorders the details mapping, so compare
+        # the pairs rather than their order.
+        (input_pairs,), _ = llm_inv.set_input_tokens.call_args
+        assert sorted(input_pairs) == [
+            ("audio", 10),
+            ("image", 20),
+            ("text", 70),
+        ]
+        (output_pairs,), _ = llm_inv.set_output_tokens.call_args
+        assert sorted(output_pairs) == [
+            ("audio", 5),
+            ("image", 10),
+            ("text", 35),
+        ]
+
+    def test_absent_token_details_forward_none(self):
+        run_id = _run_id()
+        handler, _, llm_inv = _make_handler_with_llm_invocation(run_id)
+
+        ai_msg = AIMessage(
+            content="hi",
+            usage_metadata={
+                "input_tokens": 1,
+                "output_tokens": 2,
+                "total_tokens": 3,
+            },
+        )
+        gen = ChatGeneration(
+            message=ai_msg, generation_info={"finish_reason": "stop"}
+        )
+
+        handler.on_llm_end(
+            response=LLMResult(generations=[[gen]]), run_id=run_id
+        )
+
+        llm_inv.set_input_tokens.assert_called_once_with(None)
+        llm_inv.set_output_tokens.assert_called_once_with(None)
 
 
 # ---------------------------------------------------------------------------
@@ -1728,17 +1762,13 @@ def test_extract_token_details_modalities():
             "reasoning": 10,
         },
     }
+    # Modality break-downs are forwarded to InferenceInvocation rather than
+    # flattened here, so only cache and reasoning remain.
     details = extract_token_details(usage)
     assert details == {
         "cache_write_input_tokens": 15,
         "cache_read_input_tokens": 25,
-        "text_input_tokens": 70,
-        "image_input_tokens": 20,
-        "audio_input_tokens": 10,
         "reasoning_tokens": 10,
-        "text_output_tokens": 30,
-        "image_output_tokens": 15,
-        "audio_output_tokens": 5,
     }
 
 

instrumentation/opentelemetry-instrumentation-google-genai/.changelog/674.added

@@ -0,0 +1 @@
+Record modality token usage breakdown attributes (text, image, audio) for generate_content.

instrumentation/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py

@@ -14,6 +14,7 @@
     GenerateContentConfig,
     GenerateContentConfigOrDict,
     GenerateContentResponse,
+    ModalityTokenCount,
     Tool,
     ToolUnionDict,
 )
@@ -34,6 +35,7 @@
 from opentelemetry.util.genai.types import (
     FunctionToolDefinition,
     GenericToolDefinition,
+    ModalityTokens,
     ToolDefinition,
 )
 from opentelemetry.util.types import AttributeValue
@@ -346,6 +348,17 @@ def _apply_request_attributes(
     )
 
 
+def _modality_tokens(
+    response: GenerateContentResponse, path: str
+) -> ModalityTokens | None:
+    entries: list[ModalityTokenCount] | None = _get_response_property(
+        response, path
+    )
+    if entries is None:
+        return None
+    return [(entry.modality or "", entry.token_count) for entry in entries]
+
+
 def _get_response_property(response: GenerateContentResponse, path: str):
     path_segments = path.split(".")
     current_context = response
@@ -423,6 +436,15 @@ def _apply_response_attributes(
         invocation.output_tokens = (
             invocation.output_tokens or 0
         ) + thinking_tokens
+    invocation.set_input_tokens(
+        _modality_tokens(response, "usage_metadata.prompt_tokens_details")
+    )
+    invocation.set_output_tokens(
+        _modality_tokens(response, "usage_metadata.candidates_tokens_details")
+    )
+    invocation.set_cache_read_input_tokens(
+        _modality_tokens(response, "usage_metadata.cache_tokens_details")
+    )
 
 
 def _maybe_get_tool_definitions(

instrumentation/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/interactions.py

@@ -98,6 +98,7 @@ class Stream:
     GenericPart,
     GenericToolDefinition,
     InputMessage,
+    ModalityTokens,
     OutputMessage,
     Role,
     TextPart,
@@ -152,41 +153,14 @@ def _apply_interaction_response_attributes(
     if isinstance(invocation, InferenceInvocation):
         invocation.thinking_tokens = usage.total_thought_tokens
 
-        def _set_modality_tokens(
-            entries: Any,
-            text_attr: str,
-            image_attr: str,
-            audio_attr: str,
-        ) -> None:
-            for entry in entries or []:
-                modality = _get_field(entry, "modality")
-                tokens = _get_field(entry, "tokens")
-                if modality and tokens is not None:
-                    m = str(modality).lower()
-                    if m == "text":
-                        setattr(invocation, text_attr, tokens)
-                    elif m == "image":
-                        setattr(invocation, image_attr, tokens)
-                    elif m == "audio":
-                        setattr(invocation, audio_attr, tokens)
-
-        _set_modality_tokens(
-            _get_field(usage, "input_tokens_by_modality"),
-            "text_input_tokens",
-            "image_input_tokens",
-            "audio_input_tokens",
+        invocation.set_input_tokens(
+            _modality_tokens(usage, "input_tokens_by_modality")
         )
-        _set_modality_tokens(
-            _get_field(usage, "output_tokens_by_modality"),
-            "text_output_tokens",
-            "image_output_tokens",
-            "audio_output_tokens",
+        invocation.set_output_tokens(
+            _modality_tokens(usage, "output_tokens_by_modality")
         )
-        _set_modality_tokens(
-            _get_field(usage, "cached_tokens_by_modality"),
-            "text_cache_read_input_tokens",
-            "image_cache_read_input_tokens",
-            "audio_cache_read_input_tokens",
+        invocation.set_cache_read_input_tokens(
+            _modality_tokens(usage, "cached_tokens_by_modality")
         )
 
     if telemetry_handler.should_capture_content():
@@ -195,6 +169,16 @@ def _set_modality_tokens(
         )
 
 
+def _modality_tokens(usage: Any, name: str) -> ModalityTokens | None:
+    entries = _get_field(usage, name)
+    if entries is None:
+        return None
+    return [
+        (_get_field(entry, "modality") or "", _get_field(entry, "tokens"))
+        for entry in entries
+    ]
+
+
 def _get_field(obj: Any, name: str) -> Any:
     if isinstance(obj, dict):
         return obj.get(name)